Improve creation of column and sort enums

.gitignore

@@ -11,6 +11,7 @@ __pycache__/
 # Distribution / packaging
 .Python
 env/
+.venv/
 build/
 develop-eggs/
 dist/

examples/flask_sqlalchemy/app.py

@@ -1,43 +1,46 @@
 #!/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
 
-default_query = '''
+example_query = """
 {
-  allEmployees {
+  allEmployees(sort: [NAME_ASC, ID_ASC]) {
     edges {
       node {
-        id,
-        name,
+        id
+        name
         department {
-          id,
+          id
           name
-        },
+        }
         role {
-          id,
+          id
           name
         }
       }
     }
   }
-}'''.strip()
+}
+"""
 
 
-app.add_url_rule('/graphql', view_func=GraphQLView.as_view('graphql', schema=schema, graphiql=True))
+app.add_url_rule(
+    "/graphql", view_func=GraphQLView.as_view("graphql", schema=schema, graphiql=True)
+)
 
 
 @app.teardown_appcontext
 def shutdown_session(exception=None):
     db_session.remove()
 
-if __name__ == '__main__':
+
+if __name__ == "__main__":
     init_db()
     app.run()

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,11 +1,10 @@
+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
+from graphene_sqlalchemy import SQLAlchemyConnectionField, SQLAlchemyObjectType
 
 
 class Department(SQLAlchemyObjectType):
@@ -26,18 +25,11 @@ class Meta:
         interfaces = (relay.Node, )
 
 
-SortEnumEmployee = utils.sort_enum_for_model(EmployeeModel, 'SortEnumEmployee',
-    lambda c, d: c.upper() + ('_ASC' if d else '_DESC'))
-
-
 class Query(graphene.ObjectType):
     node = relay.Node.Field()
     # Allow only single column sorting
     all_employees = SQLAlchemyConnectionField(
-        Employee,
-        sort=graphene.Argument(
-            SortEnumEmployee,
-            default_value=utils.EnumValue('id_asc', EmployeeModel.id.asc())))
+        Employee, sort=Employee.sort_argument())
     # Allows sorting over multiple columns, by default over the primary key
     all_roles = SQLAlchemyConnectionField(Role)
     # Disable sorting over this field

graphene_sqlalchemy/converter.py

@@ -7,6 +7,9 @@
                       String)
 from graphene.types.json import JSONString
 
+from .enums import enum_for_sa_enum
+from .registry import get_global_registry
+
 try:
     from sqlalchemy_utils import ChoiceType, JSONType, ScalarListType, TSVectorType
 except ImportError:
@@ -145,21 +148,15 @@ def convert_column_to_float(type, column, registry=None):
 
 @convert_sqlalchemy_type.register(types.Enum)
 def convert_enum_to_enum(type, column, registry=None):
-    enum_class = getattr(type, 'enum_class', None)
-    if enum_class:  # Check if an enum.Enum type is used
-        graphene_type = Enum.from_enum(enum_class)
-    else:  # Nope, just a list of string options
-        items = zip(type.enums, type.enums)
-        graphene_type = Enum(type.name, items)
     return Field(
-        graphene_type,
+        lambda: enum_for_sa_enum(type, registry or get_global_registry()),
         description=get_column_doc(column),
         required=not (is_column_nullable(column)),
     )
 
 
 @convert_sqlalchemy_type.register(ChoiceType)
-def convert_column_to_enum(type, column, registry=None):
+def convert_choice_to_enum(type, column, registry=None):
     name = "{}_{}".format(column.table.name, column.name).upper()
     return Enum(name, type.choices, description=get_column_doc(column))
 

graphene_sqlalchemy/enums.py

@@ -0,0 +1,203 @@
+from sqlalchemy import Column
+from sqlalchemy.types import Enum as SQLAlchemyEnumType
+
+from graphene import Argument, Enum, List
+
+from .utils import EnumValue, to_enum_value_name, to_type_name
+
+
+def _convert_sa_to_graphene_enum(sa_enum, fallback_name=None):
+    """Convert the given SQLAlchemy Enum type to a Graphene Enum type.
+
+    The name of the Graphene Enum will be determined as follows:
+    If the SQLAlchemy Enum is based on a Python Enum, use the name
+    of the Python Enum.  Otherwise, if the SQLAlchemy Enum is named,
+    use the SQL name after conversion to a type name. Otherwise, use
+    the given fallback_name or raise an error if it is empty.
+
+    The Enum value names are converted to upper case if necessary.
+    """
+    if not isinstance(sa_enum, SQLAlchemyEnumType):
+        raise TypeError(
+            "Expected sqlalchemy.types.Enum, but got: {!r}".format(sa_enum)
+        )
+    enum_class = sa_enum.enum_class
+    if enum_class:
+        if all(to_enum_value_name(key) == key for key in enum_class.__members__):
+            return Enum.from_enum(enum_class)
+        name = enum_class.__name__
+        members = [
+            (to_enum_value_name(key), value.value)
+            for key, value in enum_class.__members__.items()
+        ]
+    else:
+        sql_enum_name = sa_enum.name
+        if sql_enum_name:
+            name = to_type_name(sql_enum_name)
+        elif fallback_name:
+            name = fallback_name
+        else:
+            raise TypeError("No type name specified for {!r}".format(sa_enum))
+        members = [(to_enum_value_name(key), key) for key in sa_enum.enums]
+    return Enum(name, members)
+
+
+def enum_for_sa_enum(sa_enum, registry):
+    """Return the Graphene Enum type for the specified SQLAlchemy Enum type."""
+    if not isinstance(sa_enum, SQLAlchemyEnumType):
+        raise TypeError(
+            "Expected sqlalchemy.types.Enum, but got: {!r}".format(sa_enum)
+        )
+    enum = registry.get_graphene_enum_for_sa_enum(sa_enum)
+    if not enum:
+        enum = _convert_sa_to_graphene_enum(sa_enum)
+        registry.register_enum(sa_enum, enum)
+    return enum
+
+
+def enum_for_field(obj_type, field_name):
+    """Return the Graphene Enum type for the specified Graphene field."""
+    from .types import SQLAlchemyObjectType
+
+    if not isinstance(obj_type, type) or not issubclass(obj_type, SQLAlchemyObjectType):
+        raise TypeError(
+            "Expected SQLAlchemyObjectType, but got: {!r}".format(obj_type))
+    if not field_name or not isinstance(field_name, str):
+        raise TypeError(
+            "Expected a field name, but got: {!r}".format(field_name))
+    registry = obj_type._meta.registry
+    orm_field = registry.get_orm_field_for_graphene_field(obj_type, field_name)
+    if orm_field is None:
+        raise TypeError("Cannot get {}.{}".format(obj_type._meta.name, field_name))
+    if not isinstance(orm_field, Column):
+        raise TypeError(
+            "{}.{} does not map to model column".format(obj_type._meta.name, field_name)
+        )
+    sa_enum = orm_field.type
+    if not isinstance(sa_enum, SQLAlchemyEnumType):
+        raise TypeError(
+            "{}.{} does not map to enum column".format(obj_type._meta.name, field_name)
+        )
+    enum = registry.get_graphene_enum_for_sa_enum(sa_enum)
+    if not enum:
+        fallback_name = obj_type._meta.name + to_type_name(field_name)
+        enum = _convert_sa_to_graphene_enum(sa_enum, fallback_name)
+        registry.register_enum(sa_enum, enum)
+    return enum
+
+
+def _default_sort_enum_symbol_name(column_name, sort_asc=True):
+    return to_enum_value_name(column_name) + ("_ASC" if sort_asc else "_DESC")
+
+
+def sort_enum_for_object_type(
+    obj_type, name=None, only_fields=None, only_indexed=None, get_symbol_name=None
+):
+    """Return Graphene Enum for sorting the given SQLAlchemyObjectType.
+
+    Parameters
+    - obj_type : SQLAlchemyObjectType
+        The object type for which the sort Enum shall be generated.
+    - name : str, optional, default None
+        Name to use for the sort Enum.
+        If not provided, it will be set to the object type name + 'SortEnum'
+    - only_fields : sequence, optional, default None
+        If this is set, only fields from this sequence will be considered.
+    - only_indexed : bool, optional, default False
+        If this is set, only indexed columns will be considered.
+    - get_symbol_name : function, optional, default None
+        Function which takes the column name and a boolean indicating
+        if the sort direction is ascending, and returns the symbol name
+        for the current column and sort direction. If no such function
+        is passed, a default function will be used that creates the symbols
+        'foo_asc' and 'foo_desc' for a column with the name 'foo'.
+
+    Returns
+    - Enum
+        The Graphene Enum type
+    """
+    name = name or obj_type._meta.name + "SortEnum"
+    registry = obj_type._meta.registry
+    enum = registry.get_sort_enum_for_object_type(obj_type)
+    custom_options = dict(
+        only_fields=only_fields,
+        only_indexed=only_indexed,
+        get_symbol_name=get_symbol_name,
+    )
+    if enum:
+        if name != enum.__name__ or custom_options != enum.custom_options:
+            raise ValueError(
+                "Sort enum for {} has already been customized".format(obj_type)
+            )
+    else:
+        members = []
+        default = []
+        fields = obj_type._meta.fields
+        get_name = get_symbol_name or _default_sort_enum_symbol_name
+        for field_name in fields:
+            if only_fields and field_name not in only_fields:
+                continue
+            orm_field = registry.get_orm_field_for_graphene_field(obj_type, field_name)
+            if not isinstance(orm_field, Column):
+                continue
+            if only_indexed and not (orm_field.primary_key or orm_field.index):
+                continue
+            asc_name = get_name(orm_field.name, True)
+            asc_value = EnumValue(asc_name, orm_field.asc())
+            desc_name = get_name(orm_field.name, False)
+            desc_value = EnumValue(desc_name, orm_field.desc())
+            if orm_field.primary_key:
+                default.append(asc_value)
+            members.extend(((asc_name, asc_value), (desc_name, desc_value)))
+        enum = Enum(name, members)
+        enum.default = default  # store default as attribute
+        enum.custom_options = custom_options
+        registry.register_sort_enum(obj_type, enum)
+    return enum
+
+
+def sort_argument_for_object_type(
+    obj_type,
+    enum_name=None,
+    only_fields=None,
+    only_indexed=None,
+    get_symbol_name=None,
+    has_default=True,
+):
+    """"Returns Graphene Argument for sorting the given SQLAlchemyObjectType.
+
+    Parameters
+    - obj_type : SQLAlchemyObjectType
+        The object type for which the sort Argument shall be generated.
+    - enum_name : str, optional, default None
+        Name to use for the sort Enum.
+        If not provided, it will be set to the object type name + 'SortEnum'
+    - only_fields : sequence, optional, default None
+        If this is set, only fields from this sequence will be considered.
+    - only_indexed : bool, optional, default False
+        If this is set, only indexed columns will be considered.
+    - get_symbol_name : function, optional, default None
+        Function which takes the column name and a boolean indicating
+        if the sort direction is ascending, and returns the symbol name
+        for the current column and sort direction. If no such function
+        is passed, a default function will be used that creates the symbols
+        'foo_asc' and 'foo_desc' for a column with the name 'foo'.
+    - has_default : bool, optional, default True
+        If this is set to False, no sorting will happen when this argument is not
+        passed. Otherwise results will be sortied by the primary key(s) of the model.
+
+    Returns
+    - Enum
+        A Graphene Argument that accepts a list of sorting directions for the model.
+    """
+    enum = sort_enum_for_object_type(
+        obj_type,
+        enum_name,
+        only_fields=only_fields,
+        only_indexed=only_indexed,
+        get_symbol_name=get_symbol_name,
+    )
+    if not has_default:
+        enum.default = None
+
+    return Argument(List(enum), default_value=enum.default)