Native support for additional Type Converters

graphene_sqlalchemy/enums.py

@@ -144,9 +144,9 @@ def sort_enum_for_object_type(
             column = orm_field.columns[0]
             if only_indexed and not (column.primary_key or column.index):
                 continue
-            asc_name = get_name(column.name, True)
+            asc_name = get_name(column.key, True)
             asc_value = EnumValue(asc_name, column.asc())
-            desc_name = get_name(column.name, False)
+            desc_name = get_name(column.key, False)
             desc_value = EnumValue(desc_name, column.desc())
             if column.primary_key:
                 default.append(asc_value)

graphene_sqlalchemy/registry.py

@@ -1,7 +1,9 @@
 from collections import defaultdict
+from typing import List, Type
 
 from sqlalchemy.types import Enum as SQLAlchemyEnumType
 
+import graphene
 from graphene import Enum
 
 
@@ -13,12 +15,13 @@ def __init__(self):
         self._registry_composites = {}
         self._registry_enums = {}
         self._registry_sort_enums = {}
+        self._registry_unions = {}
 
     def register(self, obj_type):
-        from .types import SQLAlchemyObjectType
 
+        from .types import SQLAlchemyObjectType
         if not isinstance(obj_type, type) or not issubclass(
-            obj_type, SQLAlchemyObjectType
+                obj_type, SQLAlchemyObjectType
         ):
             raise TypeError(
                 "Expected SQLAlchemyObjectType, but got: {!r}".format(obj_type)
@@ -37,7 +40,7 @@ def register_orm_field(self, obj_type, field_name, orm_field):
         from .types import SQLAlchemyObjectType
 
         if not isinstance(obj_type, type) or not issubclass(
-            obj_type, SQLAlchemyObjectType
+                obj_type, SQLAlchemyObjectType
         ):
             raise TypeError(
                 "Expected SQLAlchemyObjectType, but got: {!r}".format(obj_type)
@@ -55,7 +58,7 @@ def register_composite_converter(self, composite, converter):
     def get_converter_for_composite(self, composite):
         return self._registry_composites.get(composite)
 
-    def register_enum(self, sa_enum, graphene_enum):
+    def register_enum(self, sa_enum: SQLAlchemyEnumType, graphene_enum: Enum):
         if not isinstance(sa_enum, SQLAlchemyEnumType):
             raise TypeError(
                 "Expected SQLAlchemyEnumType, but got: {!r}".format(sa_enum)
@@ -67,14 +70,14 @@ def register_enum(self, sa_enum, graphene_enum):
 
         self._registry_enums[sa_enum] = graphene_enum
 
-    def get_graphene_enum_for_sa_enum(self, sa_enum):
+    def get_graphene_enum_for_sa_enum(self, sa_enum: SQLAlchemyEnumType):
         return self._registry_enums.get(sa_enum)
 
-    def register_sort_enum(self, obj_type, sort_enum):
-        from .types import SQLAlchemyObjectType
+    def register_sort_enum(self, obj_type, sort_enum: Enum):
 
+        from .types import SQLAlchemyObjectType
         if not isinstance(obj_type, type) or not issubclass(
-            obj_type, SQLAlchemyObjectType
+                obj_type, SQLAlchemyObjectType
         ):
             raise TypeError(
                 "Expected SQLAlchemyObjectType, but got: {!r}".format(obj_type)
@@ -83,9 +86,26 @@ def register_sort_enum(self, obj_type, sort_enum):
             raise TypeError("Expected Graphene Enum, but got: {!r}".format(sort_enum))
         self._registry_sort_enums[obj_type] = sort_enum
 
-    def get_sort_enum_for_object_type(self, obj_type):
+    def get_sort_enum_for_object_type(self, obj_type: graphene.ObjectType):
         return self._registry_sort_enums.get(obj_type)
 
+    def register_union_type(self, union: graphene.Union, obj_types: List[Type[graphene.ObjectType]]):
+        if not isinstance(union, graphene.Union):
+            raise TypeError(
+                "Expected graphene.Union, but got: {!r}".format(union)
+            )
+
+        for obj_type in obj_types:
+            if not isinstance(obj_type, type(graphene.ObjectType)):
+                raise TypeError(
+                    "Expected Graphene ObjectType, but got: {!r}".format(obj_type)
+                )
+
+        self._registry_unions[frozenset(obj_types)] = union
+
+    def get_union_for_object_types(self, obj_types : List[Type[graphene.ObjectType]]):
+        return self._registry_unions.get(frozenset(obj_types))
+
 
 registry = None
 

graphene_sqlalchemy/tests/models.py

@@ -5,8 +5,8 @@
 from decimal import Decimal
 from typing import List, Optional, Tuple
 
-from sqlalchemy import (Column, Date, Enum, ForeignKey, Integer, String, Table,
-                        func, select)
+from sqlalchemy import (Column, Date, Enum, ForeignKey, Integer, Numeric,
+                        String, Table, func, select)
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import column_property, composite, mapper, relationship
@@ -228,3 +228,9 @@ def hybrid_prop_self_referential_list(self) -> List['ShoppingCart']:
     @hybrid_property
     def hybrid_prop_optional_self_referential(self) -> Optional['ShoppingCart']:
         return None
+
+
+class KeyedModel(Base):
+    __tablename__ = "test330"
+    id = Column(Integer(), primary_key=True)
+    reporter_number = Column("% reporter_number", Numeric, key="reporter_number")