diff --git a/.travis.yml b/.travis.yml index ae23cd02..544966db 100644 --- a/.travis.yml +++ b/.travis.yml @@ -23,7 +23,7 @@ before_install: install: - | if [ "$TEST_TYPE" = build ]; then - pip install --download-cache $HOME/.cache/pip/ pytest pytest-cov coveralls six pytest-django django-filter + pip install --download-cache $HOME/.cache/pip/ pytest pytest-cov coveralls six pytest-django django-filter sqlalchemy_utils pip install --download-cache $HOME/.cache/pip/ -e .[django] pip install --download-cache $HOME/.cache/pip/ -e .[sqlalchemy] pip install django==$DJANGO_VERSION diff --git a/graphene/contrib/sqlalchemy/converter.py b/graphene/contrib/sqlalchemy/converter.py index 25ed6b35..8a2b0cbe 100644 --- a/graphene/contrib/sqlalchemy/converter.py +++ b/graphene/contrib/sqlalchemy/converter.py @@ -2,7 +2,13 @@ from singledispatch import singledispatch from sqlalchemy import types from sqlalchemy.orm import interfaces +try: + from sqlalchemy_utils.types.choice import ChoiceType +except ImportError: + class ChoiceType(object): + pass +from ...core.classtypes.enum import Enum from ...core.types.scalars import ID, Boolean, Float, Int, String from .fields import ConnectionOrListField, SQLAlchemyModelField @@ -59,3 +65,9 @@ def convert_column_to_boolean(type, column): @convert_sqlalchemy_type.register(types.Numeric) def convert_column_to_float(type, column): return Float(description=column.doc) + + +@convert_sqlalchemy_type.register(ChoiceType) +def convert_column_to_enum(type, column): + name = '{}_{}'.format(column.table.name, column.name).upper() + return Enum(name, type.choices, description=column.doc) diff --git a/graphene/contrib/sqlalchemy/tests/test_converter.py b/graphene/contrib/sqlalchemy/tests/test_converter.py index e4cdaa6f..53cbcab8 100644 --- a/graphene/contrib/sqlalchemy/tests/test_converter.py +++ b/graphene/contrib/sqlalchemy/tests/test_converter.py @@ -5,7 +5,9 @@ from graphene.contrib.sqlalchemy.converter import (convert_sqlalchemy_column, convert_sqlalchemy_relationship) from graphene.contrib.sqlalchemy.fields import (ConnectionOrListField, SQLAlchemyModelField) -from sqlalchemy import Column, types +from sqlalchemy import Table, Column, types +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy_utils.types.choice import ChoiceType from .models import Article, Pet, Reporter @@ -85,6 +87,23 @@ def test_should_numeric_convert_float(): assert_column_conversion(types.Numeric(), graphene.Float) +def test_should_choice_convert_enum(): + TYPES = [ + (u'es', u'Spanish'), + (u'en', u'English') + ] + column = Column(ChoiceType(TYPES), doc='Language', name='language') + Base = declarative_base() + + Table('translatedmodel', Base.metadata, column) + graphene_type = convert_sqlalchemy_column(column) + assert issubclass(graphene_type, graphene.Enum) + assert graphene_type._meta.type_name == 'TRANSLATEDMODEL_LANGUAGE' + assert graphene_type._meta.description == 'Language' + assert graphene_type.__enum__.__members__['es'].value == 'Spanish' + assert graphene_type.__enum__.__members__['en'].value == 'English' + + def test_should_manytomany_convert_connectionorlist(): graphene_type = convert_sqlalchemy_relationship(Reporter.pets.property) assert isinstance(graphene_type, ConnectionOrListField) diff --git a/setup.py b/setup.py index b6b9d875..8dc1a871 100644 --- a/setup.py +++ b/setup.py @@ -63,6 +63,7 @@ setup( 'pytest>=2.7.2', 'pytest-django', 'sqlalchemy', + 'sqlalchemy_utils', 'mock', ], extras_require={