Register sort enums in the registry

Slightly changed the API: Creating and getting sort enums
are different functions now. When no sort enum is created,
a default one will be created automatically.

Sort enums now also use upper case symbols by default.

Author: Cito

examples/flask_sqlalchemy/app.py

@@ -1,12 +1,11 @@
 #!/usr/bin/env python
 
+from database import db_session, init_db
 from flask import Flask
+from schema import schema
 
 from flask_graphql import GraphQLView
 
-from .database import db_session, init_db
-from .schema import schema
-
 app = Flask(__name__)
 app.debug = True
 

examples/flask_sqlalchemy/database.py

@@ -14,7 +14,7 @@ def init_db():
     # import all modules here that might define models so that
     # they will be registered properly on the metadata.  Otherwise
     # you will have to import them first before calling init_db()
-    from .models import Department, Employee, Role
+    from models import Department, Employee, Role
     Base.metadata.drop_all(bind=engine)
     Base.metadata.create_all(bind=engine)
 

examples/flask_sqlalchemy/models.py

@@ -1,8 +1,7 @@
+from database import Base
 from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, func
 from sqlalchemy.orm import backref, relationship
 
-from .database import Base
-
 
 class Department(Base):
     __tablename__ = 'department'

examples/flask_sqlalchemy/schema.py

@@ -1,12 +1,12 @@
+from models import Department as DepartmentModel
+from models import Employee as EmployeeModel
+from models import Role as RoleModel
+
 import graphene
 from graphene import relay
 from graphene_sqlalchemy import (SQLAlchemyConnectionField,
                                  SQLAlchemyObjectType, utils)
 
-from .models import Department as DepartmentModel
-from .models import Employee as EmployeeModel
-from .models import Role as RoleModel
-
 
 class Department(SQLAlchemyObjectType):
     class Meta:
@@ -26,8 +26,7 @@ class Meta:
         interfaces = (relay.Node, )
 
 
-SortEnumEmployee = utils.sort_enum_for_model(EmployeeModel, 'SortEnumEmployee',
-    lambda c, d: c.upper() + ('_ASC' if d else '_DESC'))
+SortEnumEmployee = utils.get_sort_enum_for_model(EmployeeModel)
 
 
 class Query(graphene.ObjectType):

graphene_sqlalchemy/fields.py

@@ -7,7 +7,7 @@
 from graphene.relay.connection import PageInfo
 from graphql_relay.connection.arrayconnection import connection_from_list_slice
 
-from .utils import get_query, sort_argument_for_model
+from .utils import get_query, get_sort_argument_for_model
 
 
 class UnsortedSQLAlchemyConnectionField(ConnectionField):
@@ -82,7 +82,7 @@ def __init__(self, type, *args, **kwargs):
             # Let super class raise if type is not a Connection
             try:
                 model = type.Edge.node._type._meta.model
-                kwargs.setdefault("sort", sort_argument_for_model(model))
+                kwargs.setdefault("sort", get_sort_argument_for_model(model))
             except Exception:
                 raise Exception(
                     'Cannot create sort argument for {}. A model is required. Set the "sort" argument'

graphene_sqlalchemy/registry.py

@@ -1,4 +1,3 @@
-
 from collections import OrderedDict
 
 from sqlalchemy.types import Enum as SQLAlchemyEnumType
@@ -9,25 +8,76 @@
 
 
 class Registry(object):
-    def __init__(self):
+    def __init__(self, check_duplicate_registration=False):
+        self.check_duplicate_registration = check_duplicate_registration
         self._registry = {}
         self._registry_models = {}
         self._registry_composites = {}
         self._registry_enums = {}
+        self._registry_sort_params = {}
 
     def register(self, cls):
         from .types import SQLAlchemyObjectType
 
-        assert issubclass(cls, SQLAlchemyObjectType), (
-            "Only classes of type SQLAlchemyObjectType can be registered, "
-            'received "{}"'
-        ).format(cls.__name__)
-        assert cls._meta.registry == self, "Registry for a Model have to match."
-        # assert self.get_type_for_model(cls._meta.model) in [None, cls], (
-        #     'SQLAlchemy model "{}" already associated with '
-        #     'another type "{}".'
-        # ).format(cls._meta.model, self._registry[cls._meta.model])
-        self._registry[cls._meta.model] = cls
+        if not issubclass(cls, SQLAlchemyObjectType):
+            raise TypeError(
+                "Only classes of type SQLAlchemyObjectType can be registered, "
+                'received "{}"'.format(cls.__name__)
+            )
+        if cls._meta.registry != self:
+            raise TypeError("Registry for a Model have to match.")
+
+        registered_cls = (
+            self._registry.get(cls._meta.model)
+            if self.check_duplicate_registration
+            else None
+        )
+        if registered_cls:
+            if cls != registered_cls:
+                raise TypeError(
+                    "Different object types registered for the same model {}:"
+                    " tried to register {}, but {} existed already.".format(
+                        cls._meta.model, cls, registered_cls
+                    )
+                )
+        else:
+            self._registry[cls._meta.model] = cls
+
+    def register_enum(self, name, members):
+        graphene_enum = self._registry_enums.get(name)
+        if graphene_enum:
+            registered_members = {
+                key: value.value
+                for key, value in graphene_enum._meta.enum.__members__.items()
+            }
+            if members != registered_members:
+                raise TypeError(
+                    'Different enums with the same name "{}":'
+                    " tried to register {}, but {} existed already.".format(
+                        name, members, registered_members
+                    )
+                )
+        else:
+            graphene_enum = Enum(name, members)
+            self._registry_enums[name] = graphene_enum
+        return graphene_enum
+
+    def register_sort_params(self, cls, sort_params):
+        registered_sort_params = (
+            self._registry_sort_params.get(cls)
+            if self.check_duplicate_registration
+            else None
+        )
+        if registered_sort_params:
+            if registered_sort_params != sort_params:
+                raise TypeError(
+                    "Different sort args for the same model {}:"
+                    " tried to register {}, but {} existed already.".format(
+                        cls, sort_params, registered_sort_params
+                    )
+                )
+        else:
+            self._registry_sort_params[cls] = sort_params
 
     def get_type_for_model(self, model):
         return self._registry.get(model)
@@ -41,34 +91,30 @@ def get_converter_for_composite(self, composite):
     def get_type_for_enum(self, sql_type):
         if not isinstance(sql_type, SQLAlchemyEnumType):
             raise TypeError(
-                'Only sqlalchemy.Enum objects can be registered as enum, '
-                'received "{}"'.format(sql_type))
+                "Only sqlalchemy.Enum objects can be registered as enum, "
+                'received "{}"'.format(sql_type)
+            )
         enum_class = sql_type.enum_class
         if enum_class:
             name = enum_class.__name__
             members = OrderedDict(
                 (to_enum_value_name(key), value.value)
-                for key, value in enum_class.__members__.items())
+                for key, value in enum_class.__members__.items()
+            )
         else:
             name = sql_type.name
-            name = to_type_name(name) if name else 'Enum{}'.format(
-                len(self._registry_enums) + 1)
+            name = (
+                to_type_name(name)
+                if name
+                else "Enum{}".format(len(self._registry_enums) + 1)
+            )
             members = OrderedDict(
-                (to_enum_value_name(key), key) for key in sql_type.enums)
-        graphene_type = self._registry_enums.get(name)
-        if graphene_type:
-            existing_members = {
-                key: value.value for key, value
-                in graphene_type._meta.enum.__members__.items()}
-            if members != existing_members:
-                raise TypeError(
-                    'Different enums with the same name "{}":'
-                    ' tried to register {}, but {} existed already.'.format(
-                        name, members, existing_members))
-        else:
-            graphene_type = Enum(name, members)
-            self._registry_enums[name] = graphene_type
-        return graphene_type
+                (to_enum_value_name(key), key) for key in sql_type.enums
+            )
+        return self.register_enum(name, members)
+
+    def get_sort_params_for_model(self, model):
+        return self._registry_sort_params.get(model)
 
 
 registry = None

graphene_sqlalchemy/tests/test_converter.py

@@ -323,7 +323,7 @@ def test_should_postgresql_hstore_convert():
 
 
 def test_should_composite_convert():
-    class CompositeClass(object):
+    class CompositeClass:
         def __init__(self, col1, col2):
             self.col1 = col1
             self.col2 = col2

graphene_sqlalchemy/tests/test_fields.py

@@ -4,7 +4,7 @@
 
 from ..fields import SQLAlchemyConnectionField
 from ..types import SQLAlchemyObjectType
-from ..utils import sort_argument_for_model
+from ..utils import get_sort_argument_for_model
 from .models import Editor
 from .models import Pet as PetModel
 
@@ -22,7 +22,7 @@ class Meta:
 def test_sort_added_by_default():
     arg = SQLAlchemyConnectionField(PetConn)
     assert "sort" in arg.args
-    assert arg.args["sort"] == sort_argument_for_model(PetModel)
+    assert arg.args["sort"] == get_sort_argument_for_model(PetModel)
 
 
 def test_sort_can_be_removed():
@@ -31,8 +31,8 @@ def test_sort_can_be_removed():
 
 
 def test_custom_sort():
-    arg = SQLAlchemyConnectionField(PetConn, sort=sort_argument_for_model(Editor))
-    assert arg.args["sort"] == sort_argument_for_model(Editor)
+    arg = SQLAlchemyConnectionField(PetConn, sort=get_sort_argument_for_model(Editor))
+    assert arg.args["sort"] == get_sort_argument_for_model(Editor)
 
 
 def test_init_raises():

graphene_sqlalchemy/tests/test_query.py

@@ -8,7 +8,7 @@
 from ..fields import SQLAlchemyConnectionField
 from ..registry import reset_global_registry
 from ..types import SQLAlchemyObjectType
-from ..utils import sort_argument_for_model, sort_enum_for_model
+from ..utils import get_sort_argument_for_model, get_sort_enum_for_model
 from .models import Article, Base, Editor, HairKind, Pet, Reporter
 
 db = create_engine("sqlite://")  # use in-memory database
@@ -478,10 +478,10 @@ class Query(graphene.ObjectType):
         multipleSort = SQLAlchemyConnectionField(PetConnection)
         descSort = SQLAlchemyConnectionField(PetConnection)
         singleColumnSort = SQLAlchemyConnectionField(
-            PetConnection, sort=graphene.Argument(sort_enum_for_model(Pet))
+            PetConnection, sort=graphene.Argument(get_sort_enum_for_model(Pet))
         )
         noDefaultSort = SQLAlchemyConnectionField(
-            PetConnection, sort=sort_argument_for_model(Pet, False)
+            PetConnection, sort=get_sort_argument_for_model(Pet, False)
         )
         noSort = SQLAlchemyConnectionField(PetConnection, sort=None)
 
@@ -494,36 +494,36 @@ class Query(graphene.ObjectType):
                     }
                 }
             }
-            nameSort(sort: name_asc){
+            nameSort(sort: NAME_ASC){
                 edges{
                     node{
                         name
                     }
                 }
             }
-            multipleSort(sort: [pet_kind_asc, name_desc]){
+            multipleSort(sort: [PET_KIND_ASC, NAME_DESC]){
                 edges{
                     node{
                         name
                         petKind
                     }
                 }
             }
-            descSort(sort: [name_desc]){
+            descSort(sort: [NAME_DESC]){
                 edges{
                     node{
                         name
                     }
                 }
             }
-            singleColumnSort(sort: name_desc){
+            singleColumnSort(sort: NAME_DESC){
                 edges{
                     node{
                         name
                     }
                 }
             }
-            noDefaultSort(sort: name_asc){
+            noDefaultSort(sort: NAME_ASC){
                 edges{
                     node{
                         name