Improve auto creation of Graphene Enums.

The created Graphene Enums are now registered and reused,
because their names must be unique in a GraphQL schema.
Also the naming conventions for Enum type names (CamelCase)
and options (UPPERCASE) are applied when creating them.

Author: Cito

graphene_sqlalchemy/converter.py

@@ -131,12 +131,12 @@ def convert_column_to_float(type, column, registry=None):
 
 
 @convert_sqlalchemy_type.register(types.Enum)
-def convert_enum_to_enum(type, column, registry=None):
-    try:
-        items = type.enum_class.__members__.items()
-    except AttributeError:
-        items = zip(type.enums, type.enums)
-    return Field(Enum(type.name, items),
+def convert_column_to_enum(type, column, registry=None):
+    if registry is None:
+        from .registry import get_global_registry
+        registry = get_global_registry()
+    type = registry.get_type_for_enum(type)
+    return Field(type,
                  description=get_column_doc(column),
                  required=not(is_column_nullable(column)))
 

graphene_sqlalchemy/registry.py

@@ -1,9 +1,16 @@
+from graphene import Enum
+
+from sqlalchemy.types import Enum as SQLAlchemyEnumType
+
+from .utils import to_type_name
+
+
 class Registry(object):
 
     def __init__(self):
         self._registry = {}
-        self._registry_models = {}
         self._registry_composites = {}
+        self._registry_enums = {}
 
     def register(self, cls):
         from .types import SQLAlchemyObjectType
@@ -27,6 +34,35 @@ def register_composite_converter(self, composite, converter):
     def get_converter_for_composite(self, composite):
         return self._registry_composites.get(composite)
 
+    def get_type_for_enum(self, sql_type):
+        assert isinstance(sql_type, SQLAlchemyEnumType), (
+            'Only sqlalchemy.Enum objects can be registered as enum, '
+            'received "{}"'
+        ).format(sql_type)
+        if sql_type.enum_class:
+            name = sql_type.enum_class.__name__
+            items = [(key.upper(), value.value)
+                     for key, value in sql_type.enum_class.__members__.items()]
+        else:
+            name = to_type_name(sql_type.name)
+            if not name:
+                name = 'Enum{}'.format(len(self._registry_enums) + 1)
+            items = [(key.upper(), key) for key in sql_type.enums]
+        if name:
+            gql_type = self._registry_enums.get(name)
+            if gql_type:
+                if dict(items) != {
+                        key: value.value for key, value
+                            in gql_type._meta.enum.__members__.items()}:
+                    raise TypeError(
+                        'Different enums with the same name {}'.format(name))
+        else:
+            name = 'Enum{}'.format(len(self._registry_enums) + 1)
+            gql_type = None
+        if not gql_type:
+            gql_type = Enum(name, items)
+            self._registry_enums[name] = gql_type
+        return gql_type
 
 registry = None
 

graphene_sqlalchemy/tests/models.py

@@ -1,7 +1,5 @@
 from __future__ import absolute_import
 
-import enum
-
 from sqlalchemy import Column, Date, Enum, ForeignKey, Integer, String, Table
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import mapper, relationship
@@ -13,6 +11,9 @@
                           Column('reporter_id', Integer, ForeignKey('reporters.id')))
 
 
+PetKind = Enum('cat', 'dog', name='pet_kind_enum')
+
+
 class Editor(Base):
     __tablename__ = 'editors'
     editor_id = Column(Integer(), primary_key=True)
@@ -23,8 +24,7 @@ class Pet(Base):
     __tablename__ = 'pets'
     id = Column(Integer(), primary_key=True)
     name = Column(String(30))
-    pet_kind = Column(Enum('cat', 'dog', name='pet_kind'), nullable=False)
-    reporter_id = Column(Integer(), ForeignKey('reporters.id'))
+    pet_kind = Column(PetKind, nullable=False)
 
 
 class Reporter(Base):
@@ -33,6 +33,7 @@ class Reporter(Base):
     first_name = Column(String(30))
     last_name = Column(String(30))
     email = Column(String())
+    favorite_pet_kind = Column(PetKind)
     pets = relationship('Pet', secondary=association_table, backref='reporters')
     articles = relationship('Article', backref='reporter')
     favorite_article = relationship("Article", uselist=False)

graphene_sqlalchemy/tests/test_converter.py

@@ -81,16 +81,21 @@ def test_should_unicodetext_convert_string():
 
 def test_should_enum_convert_enum():
     field = assert_column_conversion(
-        types.Enum(enum.Enum('one', 'two')), graphene.Field)
+        types.Enum(enum.Enum('TwoNumbersPyEnum', 'one two')), graphene.Field)
     field_type = field.type()
+    assert field_type.__class__.__name__ == 'TwoNumbersPyEnum'
     assert isinstance(field_type, graphene.Enum)
-    assert hasattr(field_type, 'two')
+    assert hasattr(field_type, 'ONE')
+    assert not hasattr(field_type, 'one')
+    assert hasattr(field_type, 'TWO')
     field = assert_column_conversion(
-        types.Enum('one', 'two', name='two_numbers'), graphene.Field)
+        types.Enum('one', 'two', name='two_numbers_db_enum'), graphene.Field)
     field_type = field.type()
-    assert field_type.__class__.__name__ == 'two_numbers'
+    assert field_type.__class__.__name__ == 'TwoNumbersDbEnum'
     assert isinstance(field_type, graphene.Enum)
-    assert hasattr(field_type, 'two')
+    assert hasattr(field_type, 'ONE')
+    assert not hasattr(field_type, 'one')
+    assert hasattr(field_type, 'TWO')
 
 
 def test_should_small_integer_convert_int():
@@ -262,11 +267,11 @@ def test_should_postgresql_uuid_convert():
 
 def test_should_postgresql_enum_convert():
     field = assert_column_conversion(postgresql.ENUM(
-        enum.Enum('one', 'two'), name='two_numbers'), graphene.Field)
+        enum.Enum('TwoNumbers', 'one two')), graphene.Field)
     field_type = field.type()
-    assert field_type.__class__.__name__ == 'two_numbers'
+    assert field_type.__class__.__name__ == 'TwoNumbers'
     assert isinstance(field_type, graphene.Enum)
-    assert hasattr(field_type, 'two')
+    assert hasattr(field_type, 'TWO')
 
 
 def test_should_postgresql_array_convert():

graphene_sqlalchemy/tests/test_query.py

@@ -13,6 +13,16 @@
 db = create_engine('sqlite:///test_sqlalchemy.sqlite3')
 
 
+def normalize(value):
+    """Convert nested ordered dicts ot normal dicts for better comparison."""
+    if isinstance(value, dict):
+        return {k: normalize(v) for k, v in value.items()}
+    elif isinstance(value, list):
+        return [normalize(v) for v in value]
+    else:
+        return value
+
+
 @pytest.yield_fixture(scope='function')
 def session():
     reset_global_registry()
@@ -33,23 +43,34 @@ def session():
 
 
 def setup_fixtures(session):
-    pet = Pet(name='Lassie', pet_kind='dog')
-    session.add(pet)
-    reporter = Reporter(first_name='ABA', last_name='X')
+    reporter = Reporter(
+        first_name='John', last_name='Doe', favorite_pet_kind='cat')
     session.add(reporter)
-    reporter2 = Reporter(first_name='ABO', last_name='Y')
-    session.add(reporter2)
+    pet = Pet(name='Garfield', pet_kind='cat')
+    session.add(pet)
+    pet.reporters.append(reporter)
     article = Article(headline='Hi!')
     article.reporter = reporter
     session.add(article)
-    editor = Editor(name="John")
+    reporter = Reporter(
+        first_name='Jane', last_name='Roe', favorite_pet_kind='dog')
+    session.add(reporter)
+    pet = Pet(name='Lassie', pet_kind='dog')
+    pet.reporters.append(reporter)
+    session.add(pet)
+    editor = Editor(name="Jack")
     session.add(editor)
     session.commit()
 
 
 def test_should_query_well(session):
     setup_fixtures(session)
 
+    class PetType(SQLAlchemyObjectType):
+
+        class Meta:
+            model = Pet
+
     class ReporterType(SQLAlchemyObjectType):
 
         class Meta:
@@ -58,75 +79,68 @@ class Meta:
     class Query(graphene.ObjectType):
         reporter = graphene.Field(ReporterType)
         reporters = graphene.List(ReporterType)
+        pets = graphene.List(PetType, kind=graphene.Argument(
+            PetType._meta.fields['pet_kind'].type))
 
-        def resolve_reporter(self, *args, **kwargs):
+        def resolve_reporter(self, _info):
             return session.query(Reporter).first()
 
-        def resolve_reporters(self, *args, **kwargs):
+        def resolve_reporters(self, _info):
             return session.query(Reporter)
 
+        def resolve_pets(self, _info, kind):
+            query = session.query(Pet)
+            if kind:
+                query = query.filter_by(pet_kind=kind)
+            return query
+
     query = '''
         query ReporterQuery {
           reporter {
             firstName,
             lastName,
-            email
+            email,
+            favoritePetKind,
+            pets {
+              name
+              petKind
+            }
           }
           reporters {
             firstName
           }
+          pets(kind: DOG) {
+            name
+            petKind
+          }
         }
     '''
     expected = {
         'reporter': {
-            'firstName': 'ABA',
-            'lastName': 'X',
-            'email': None
+            'firstName': 'John',
+            'lastName': 'Doe',
+            'email': None,
+            'favoritePetKind': 'CAT',
+            'pets': [{
+                'name': 'Garfield',
+                'petKind': 'CAT'
+            }]
         },
         'reporters': [{
-            'firstName': 'ABA',
+            'firstName': 'John',
         }, {
-            'firstName': 'ABO',
-        }]
-    }
-    schema = graphene.Schema(query=Query)
-    result = schema.execute(query)
-    assert not result.errors
-    assert result.data == expected
-
-
-def test_should_query_enums(session):
-    setup_fixtures(session)
-
-    class PetType(SQLAlchemyObjectType):
-
-        class Meta:
-            model = Pet
-
-    class Query(graphene.ObjectType):
-        pet = graphene.Field(PetType)
-
-        def resolve_pet(self, *args, **kwargs):
-            return session.query(Pet).first()
-
-    query = '''
-        query PetQuery {
-          pet {
-            name,
-            petKind
-          }
-        }
-    '''
-    expected = {
-        'pet': {
+            'firstName': 'Jane',
+        }],
+        'pets': [{
             'name': 'Lassie',
-            'petKind': 'dog'
-        }
+            'petKind': 'DOG'
+        }]
     }
     schema = graphene.Schema(query=Query)
     result = schema.execute(query)
     assert not result.errors
-    assert result.data == expected, result.data
+    result = normalize(result.data)
+    assert result == expected
 
 
 def test_should_node(session):
@@ -158,10 +172,10 @@ class Query(graphene.ObjectType):
         article = graphene.Field(ArticleNode)
         all_articles = SQLAlchemyConnectionField(ArticleNode)
 
-        def resolve_reporter(self, *args, **kwargs):
+        def resolve_reporter(self, _info):
             return session.query(Reporter).first()
 
-        def resolve_article(self, *args, **kwargs):
+        def resolve_article(self, _info):
             return session.query(Article).first()
 
     query = '''
@@ -200,8 +214,8 @@ def resolve_article(self, *args, **kwargs):
     expected = {
         'reporter': {
             'id': 'UmVwb3J0ZXJOb2RlOjE=',
-            'firstName': 'ABA',
-            'lastName': 'X',
+            'firstName': 'John',
+            'lastName': 'Doe',
             'email': None,
             'articles': {
                 'edges': [{
@@ -226,7 +240,8 @@ def resolve_article(self, *args, **kwargs):
     schema = graphene.Schema(query=Query)
     result = schema.execute(query, context_value={'session': session})
     assert not result.errors
-    assert result.data == expected
+    result = normalize(result.data)
+    assert result == expected
 
 
 def test_should_custom_identifier(session):
@@ -264,12 +279,12 @@ class Query(graphene.ObjectType):
             'edges': [{
                 'node': {
                     'id': 'RWRpdG9yTm9kZTox',
-                    'name': 'John'
+                    'name': 'Jack'
                 }
             }]
         },
         'node': {
-            'name': 'John'
+            'name': 'Jack'
         }
     }
 
@@ -355,7 +370,7 @@ class Mutation(graphene.ObjectType):
                 'headline': 'My Article',
                 'reporter': {
                     'id': 'UmVwb3J0ZXJOb2RlOjE=',
-                    'firstName': 'ABA'
+                    'firstName': 'John'
                 }
             }
         },
@@ -364,4 +379,6 @@ class Mutation(graphene.ObjectType):
     schema = graphene.Schema(query=Query, mutation=Mutation)
     result = schema.execute(query, context_value={'session': session})
     assert not result.errors
-    assert result.data == expected
+    result = normalize(result.data)
+    assert result == expected
+

graphene_sqlalchemy/tests/test_schema.py

@@ -34,6 +34,7 @@ class Meta:
         'first_name',
         'last_name',
         'email',
+        'favorite_pet_kind',
         'pets',
         'articles',
         'favorite_article']

graphene_sqlalchemy/tests/test_types.py

@@ -50,6 +50,7 @@ def test_objecttype_registered():
         'first_name',
         'last_name',
         'email',
+        'favorite_pet_kind',
         'pets',
         'articles',
         'favorite_article']
@@ -113,6 +114,7 @@ def test_custom_objecttype_registered():
         'first_name',
         'last_name',
         'email',
+        'favorite_pet_kind',
         'pets',
         'articles',
         'favorite_article']

graphene_sqlalchemy/tests/test_utils.py

@@ -1,6 +1,6 @@
 from graphene import ObjectType, Schema, String
 
-from ..utils import get_session
+from ..utils import get_session, to_type_name
 
 
 def test_get_session():
@@ -22,3 +22,8 @@ def resolve_x(self, info):
     result = schema.execute(query, context_value={'session': session})
     assert not result.errors
     assert result.data['x'] == session
+
+
+def test_get_enum_name():
+    assert to_type_name('make_camel_case') == 'MakeCamelCase'
+    assert to_type_name('AlreadyCamelCase') == 'AlreadyCamelCase'

graphene_sqlalchemy/utils.py

@@ -34,3 +34,8 @@ def is_mapped_instance(cls):
         return False
     else:
         return True
+
+
+def to_type_name(name):
+    return ''.join(part[:1].upper() + part[1:] for part in name.split('_'))
+