From c0bd2c2bca0699b698ca66817a0c1c25a752eab6 Mon Sep 17 00:00:00 2001
From: Konsta Vesterinen <konsta.vesterinen@gmail.com>
Date: Fri, 5 Jun 2015 14:29:48 +0300
Subject: [PATCH] Added get_type utility function

---
 CHANGES.rst                            |  6 ++++
 docs/orm_helpers.rst                   |  6 ++++
 sqlalchemy_utils/__init__.py           |  3 +-
 sqlalchemy_utils/functions/__init__.py |  1 +
 sqlalchemy_utils/functions/orm.py      | 50 +++++++++++++++++++++++++-
 tests/functions/test_get_type.py       | 43 ++++++++++++++++++++++
 6 files changed, 107 insertions(+), 2 deletions(-)
 create mode 100644 tests/functions/test_get_type.py

diff --git a/CHANGES.rst b/CHANGES.rst
index a5ae7da..f3a4d15 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -4,6 +4,12 @@ Changelog
 Here you can see the full list of changes between each SQLAlchemy-Utils release.
 
 
+0.30.9 (2015-06-05)
+^^^^^^^^^^^^^^^^^^^
+
+- Added get_type utility function
+
+
 0.30.8 (2015-06-05)
 ^^^^^^^^^^^^^^^^^^^
 
diff --git a/docs/orm_helpers.rst b/docs/orm_helpers.rst
index 844a07a..eac570e 100644
--- a/docs/orm_helpers.rst
+++ b/docs/orm_helpers.rst
@@ -70,6 +70,12 @@ get_tables
 .. autofunction:: get_tables
 
 
+get_type
+--------
+
+.. autofunction:: get_type
+
+
 has_changes
 -----------
 
diff --git a/sqlalchemy_utils/__init__.py b/sqlalchemy_utils/__init__.py
index 5bcae1c..8acf281 100644
--- a/sqlalchemy_utils/__init__.py
+++ b/sqlalchemy_utils/__init__.py
@@ -28,6 +28,7 @@ from .functions import (  # noqa
     get_query_entities,
     get_referencing_foreign_keys,
     get_tables,
+    get_type,
     group_foreign_keys,
     has_changes,
     has_index,
@@ -91,4 +92,4 @@ from .types import (  # noqa
     WeekDaysType
 )
 
-__version__ = '0.30.8'
+__version__ = '0.30.9'
diff --git a/sqlalchemy_utils/functions/__init__.py b/sqlalchemy_utils/functions/__init__.py
index a02df98..a8409df 100644
--- a/sqlalchemy_utils/functions/__init__.py
+++ b/sqlalchemy_utils/functions/__init__.py
@@ -29,6 +29,7 @@ from .orm import (  # noqa
     get_primary_keys,
     get_query_entities,
     get_tables,
+    get_type,
     getdotattr,
     has_changes,
     identity,
diff --git a/sqlalchemy_utils/functions/orm.py b/sqlalchemy_utils/functions/orm.py
index 4b36d1c..d90c272 100644
--- a/sqlalchemy_utils/functions/orm.py
+++ b/sqlalchemy_utils/functions/orm.py
@@ -14,7 +14,7 @@ from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import mapperlib
 from sqlalchemy.orm.attributes import InstrumentedAttribute
 from sqlalchemy.orm.exc import UnmappedInstanceError
-from sqlalchemy.orm.properties import ColumnProperty
+from sqlalchemy.orm.properties import ColumnProperty, RelationshipProperty
 from sqlalchemy.orm.query import _ColumnEntity
 from sqlalchemy.orm.session import object_session
 from sqlalchemy.orm.util import AliasedInsp
@@ -105,6 +105,54 @@ def get_class_by_table(base, table, data=None):
     return None
 
 
+def get_type(expr):
+    """
+    Return the associated type with given Column, InstrumentedAttribute,
+    ColumnProperty, RelationshipProperty or other similar SQLAlchemy construct.
+
+    For constructs wrapping columns this is the column type. For relationships
+    this function returns the relationship mapper class.
+
+    :param expr:
+        SQLAlchemy Column, InstrumentedAttribute, ColumnProperty or other
+        similar SA construct.
+
+    ::
+
+        class User(Base):
+            __tablename__ = 'user'
+            id = sa.Column(sa.Integer, primary_key=True)
+            name = sa.Column(sa.String)
+
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+            author_id = sa.Column(sa.Integer, sa.ForeignKey(User.id))
+            author = sa.orm.relationship(User)
+
+
+        get_type(User.__table__.c.name)  # sa.String()
+        get_type(User.name)  # sa.String()
+        get_type(User.name.property)  # sa.String()
+
+        get_type(Article.author)  # User
+
+
+    .. versionadded: 0.30.9
+    """
+    if hasattr(expr, 'type'):
+        return expr.type
+    elif isinstance(expr, InstrumentedAttribute):
+        expr = expr.property
+
+    if isinstance(expr, ColumnProperty):
+        return expr.columns[0].type
+    elif isinstance(expr, RelationshipProperty):
+        return expr.mapper.class_
+    raise TypeError("Couldn't inspect type.")
+
+
 def get_column_key(model, column):
     """
     Return the key for given column in given model.
diff --git a/tests/functions/test_get_type.py b/tests/functions/test_get_type.py
new file mode 100644
index 0000000..0287325
--- /dev/null
+++ b/tests/functions/test_get_type.py
@@ -0,0 +1,43 @@
+import sqlalchemy as sa
+from sqlalchemy.ext.declarative import declarative_base
+
+from sqlalchemy_utils import get_type
+from tests import TestCase
+
+
+class TestGetType(object):
+    def setup_method(self, method):
+        Base = declarative_base()
+
+        class User(Base):
+            __tablename__ = 'user'
+            id = sa.Column(sa.Integer, primary_key=True)
+
+        class Article(Base):
+            __tablename__ = 'article'
+            id = sa.Column(sa.Integer, primary_key=True)
+
+            author_id = sa.Column(sa.Integer, sa.ForeignKey(User.id))
+            author = sa.orm.relationship(User)
+
+            some_property = sa.orm.column_property(
+                sa.func.coalesce(id, 1)
+            )
+
+        self.Article = Article
+        self.User = User
+
+    def test_instrumented_attribute(self):
+        assert isinstance(get_type(self.Article.id), sa.Integer)
+
+    def test_column_property(self):
+        assert isinstance(get_type(self.Article.id.property), sa.Integer)
+
+    def test_column(self):
+        assert isinstance(get_type(self.Article.__table__.c.id), sa.Integer)
+
+    def test_calculated_column_property(self):
+        assert isinstance(get_type(self.Article.some_property), sa.Integer)
+
+    def test_relationship_property(self):
+        assert get_type(self.Article.author) == self.User