Merge pull request #5 from cjw296/pg-ranges

Support for Postgres range types.

4 files changed, 246 insertions, 2 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/__init__.py b/lib/sqlalchemy/dialects/postgresql/__init__.py
index d0f785bdd..408b67846 100644
--- a/lib/sqlalchemy/dialects/postgresql/__init__.py
+++ b/lib/sqlalchemy/dialects/postgresql/__init__.py
@@ -12,12 +12,16 @@ from .base import \
     INTEGER, BIGINT, SMALLINT, VARCHAR, CHAR, TEXT, NUMERIC, FLOAT, REAL, \
     INET, CIDR, UUID, BIT, MACADDR, DOUBLE_PRECISION, TIMESTAMP, TIME, \
     DATE, BYTEA, BOOLEAN, INTERVAL, ARRAY, ENUM, dialect, array, Any, All
+from .constraints import ExcludeConstraint
 from .hstore import HSTORE, hstore
+from .ranges import INT4RANGE, INT8RANGE, NUMRANGE, DATERANGE, TSRANGE, \
+    TSTZRANGE
 
 __all__ = (
     'INTEGER', 'BIGINT', 'SMALLINT', 'VARCHAR', 'CHAR', 'TEXT', 'NUMERIC',
     'FLOAT', 'REAL', 'INET', 'CIDR', 'UUID', 'BIT', 'MACADDR',
     'DOUBLE_PRECISION', 'TIMESTAMP', 'TIME', 'DATE', 'BYTEA', 'BOOLEAN',
     'INTERVAL', 'ARRAY', 'ENUM', 'dialect', 'Any', 'All', 'array', 'HSTORE',
-    'hstore'
+    'hstore', 'INT4RANGE', 'INT8RANGE', 'NUMRANGE', 'DATERANGE',
+    'TSRANGE', 'TSTZRANGE'
 )
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index dc6e30b81..16ace0583 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -443,7 +443,7 @@ class array(expression.Tuple):
 
     An instance of :class:`.array` will always have the datatype
     :class:`.ARRAY`.  The "inner" type of the array is inferred from
-    the values present, unless the "type_" keyword argument is passed::
+    the values present, unless the ``type_`` keyword argument is passed::
 
         array(['foo', 'bar'], type_=CHAR)
 
@@ -1141,6 +1141,22 @@ class PGDDLCompiler(compiler.DDLCompiler):
             text += " WHERE " + where_compiled
         return text
 
+    def visit_exclude_constraint(self, constraint):
+        text = ""
+        if constraint.name is not None:
+            text += "CONSTRAINT %s " % \
+                    self.preparer.format_constraint(constraint)
+        elements = []
+        for c in constraint.columns:
+            op = constraint.operators[c.name]
+            elements.append(self.preparer.quote(c.name, c.quote)+' WITH '+op)
+        text += "EXCLUDE USING %s (%s)" % (constraint.using, ', '.join(elements))
+        if constraint.where is not None:
+            sqltext = sql_util.expression_as_ddl(constraint.where)
+            text += ' WHERE (%s)' % self.sql_compiler.process(sqltext)
+        text += self.define_constraint_deferrability(constraint)
+        return text
+
 
 class PGTypeCompiler(compiler.GenericTypeCompiler):
     def visit_INET(self, type_):
@@ -1167,6 +1183,24 @@ class PGTypeCompiler(compiler.GenericTypeCompiler):
     def visit_HSTORE(self, type_):
         return "HSTORE"
 
+    def visit_INT4RANGE(self, type_):
+        return "INT4RANGE"
+
+    def visit_INT8RANGE(self, type_):
+        return "INT8RANGE"
+
+    def visit_NUMRANGE(self, type_):
+        return "NUMRANGE"
+
+    def visit_DATERANGE(self, type_):
+        return "DATERANGE"
+
+    def visit_TSRANGE(self, type_):
+        return "TSRANGE"
+
+    def visit_TSTZRANGE(self, type_):
+        return "TSTZRANGE"
+
     def visit_datetime(self, type_):
         return self.visit_TIMESTAMP(type_)
 
diff --git a/lib/sqlalchemy/dialects/postgresql/constraints.py b/lib/sqlalchemy/dialects/postgresql/constraints.py
new file mode 100644
index 000000000..5b8bbe643
--- /dev/null
+++ b/lib/sqlalchemy/dialects/postgresql/constraints.py
@@ -0,0 +1,73 @@
+# Copyright (C) 2013 the SQLAlchemy authors and contributors <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: http://www.opensource.org/licenses/mit-license.php
+from sqlalchemy.schema import ColumnCollectionConstraint
+from sqlalchemy.sql import expression
+
+class ExcludeConstraint(ColumnCollectionConstraint):
+    """A table-level EXCLUDE constraint.
+
+    Defines an EXCLUDE constraint as described in the `postgres
+    documentation`__.
+
+    __ http://www.postgresql.org/docs/9.0/static/sql-createtable.html#SQL-CREATETABLE-EXCLUDE
+    """
+
+    __visit_name__ = 'exclude_constraint'
+
+    where = None
+
+    def __init__(self, *elements, **kw):
+        """
+        :param \*elements:
+          A sequence of two tuples of the form ``(column, operator)`` where
+          column must be a column name or Column object and operator must
+          be a string containing the operator to use.
+
+        :param name:
+          Optional, the in-database name of this constraint.
+
+        :param deferrable:
+          Optional bool.  If set, emit DEFERRABLE or NOT DEFERRABLE when
+          issuing DDL for this constraint.
+
+        :param initially:
+          Optional string.  If set, emit INITIALLY <value> when issuing DDL
+          for this constraint.
+
+        :param using:
+          Optional string.  If set, emit USING <index_method> when issuing DDL
+          for this constraint. Defaults to 'gist'.
+          
+        :param where:
+          Optional string.  If set, emit WHERE <predicate> when issuing DDL
+          for this constraint.
+
+        """
+        ColumnCollectionConstraint.__init__(
+            self,
+            *[col for col, op in elements],
+            name=kw.get('name'),
+            deferrable=kw.get('deferrable'),
+            initially=kw.get('initially')
+            )
+        self.operators = {}
+        for col_or_string, op in elements:
+            name = getattr(col_or_string, 'name', col_or_string)
+            self.operators[name] = op
+        self.using = kw.get('using', 'gist')
+        where = kw.get('where')
+        if where:
+            self.where =  expression._literal_as_text(where)
+            
+    def copy(self, **kw):
+        elements = [(col, self.operators[col])
+                    for col in self.columns.keys()]
+        c = self.__class__(*elements,
+                            name=self.name,
+                            deferrable=self.deferrable,
+                            initially=self.initially)
+        c.dispatch._update(self.dispatch)
+        return c
+
diff --git a/lib/sqlalchemy/dialects/postgresql/ranges.py b/lib/sqlalchemy/dialects/postgresql/ranges.py
new file mode 100644
index 000000000..e7ab1d5b5
--- /dev/null
+++ b/lib/sqlalchemy/dialects/postgresql/ranges.py
@@ -0,0 +1,133 @@
+# Copyright (C) 2013 the SQLAlchemy authors and contributors <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: http://www.opensource.org/licenses/mit-license.php
+
+from .base import ischema_names
+from ... import types as sqltypes
+
+__all__ = ('INT4RANGE', 'INT8RANGE', 'NUMRANGE')
+
+class RangeOperators(object):
+    """
+    This mixin provides functionality for the Range Operators
+    listed in Table 9-44 of the `postgres documentation`__ for Range
+    Functions and Operators. It is used by all the range types
+    provided in the ``postgres`` dialect and can likely be used for
+    any range types you create yourself.
+
+    __ http://www.postgresql.org/docs/devel/static/functions-range.html
+
+    No extra support is provided for the Range Functions listed in
+    Table 9-45 of the postgres documentation. For these, the normal
+    :func:`~sqlalchemy.sql.expression.func` object should be used.
+    """
+
+    class comparator_factory(sqltypes.Concatenable.Comparator):
+        """Define comparison operations for range types."""
+
+        def __ne__(self, other):
+            "Boolean expression. Returns true if two ranges are not equal" 
+            return self.expr.op('<>')(other)
+
+        def contains(self, other, **kw):
+            """Boolean expression. Returns true if the right hand operand,
+            which can be an element or a range, is contained within the
+            column.
+            """
+            return self.expr.op('@>')(other)
+
+        def contained_by(self, other):
+            """Boolean expression. Returns true if the column is contained
+            within the right hand operand.
+            """
+            return self.expr.op('<@')(other)
+
+        def overlaps(self, other):
+            """Boolean expression. Returns true if the column overlaps
+            (has points in common with) the right hand operand.
+            """
+            return self.expr.op('&&')(other)
+
+        def strictly_left_of(self, other):
+            """Boolean expression. Returns true if the column is strictly
+            left of the right hand operand.
+            """
+            return self.expr.op('<<')(other)
+            
+        __lshift__ = strictly_left_of
+
+        def strictly_right_of(self, other):
+            """Boolean expression. Returns true if the column is strictly
+            right of the right hand operand.
+            """
+            return self.expr.op('>>')(other)
+            
+        __rshift__ = strictly_right_of
+        
+        def not_extend_right_of(self, other):
+            """Boolean expression. Returns true if the range in the column
+            does not extend right of the range in the operand.
+            """
+            return self.expr.op('&<')(other)
+        
+        def not_extend_left_of(self, other):
+            """Boolean expression. Returns true if the range in the column
+            does not extend left of the range in the operand.
+            """
+            return self.expr.op('&>')(other)
+
+        def adjacent_to(self, other):
+            """Boolean expression. Returns true if the range in the column
+            is adjacent to the range in the operand.
+            """
+            return self.expr.op('-|-')(other)
+
+        def __add__(self, other):
+            """Range expression. Returns the union of the two ranges.
+            Will raise an exception if the resulting range is not
+            contigous.
+            """
+            return self.expr.op('+')(other)
+
+class INT4RANGE(RangeOperators, sqltypes.TypeEngine):
+    "Represent the Postgresql INT4RANGE type."
+    
+    __visit_name__ = 'INT4RANGE'
+
+ischema_names['int4range'] = INT4RANGE
+
+class INT8RANGE(RangeOperators, sqltypes.TypeEngine):
+    "Represent the Postgresql INT8RANGE type."
+    
+    __visit_name__ = 'INT8RANGE'
+
+ischema_names['int8range'] = INT8RANGE
+
+class NUMRANGE(RangeOperators, sqltypes.TypeEngine):
+    "Represent the Postgresql NUMRANGE type."
+    
+    __visit_name__ = 'NUMRANGE'
+
+ischema_names['numrange'] = NUMRANGE
+
+class DATERANGE(RangeOperators, sqltypes.TypeEngine):
+    "Represent the Postgresql DATERANGE type."
+    
+    __visit_name__ = 'DATERANGE'
+
+ischema_names['daterange'] = DATERANGE
+
+class TSRANGE(RangeOperators, sqltypes.TypeEngine):
+    "Represent the Postgresql TSRANGE type."
+    
+    __visit_name__ = 'TSRANGE'
+
+ischema_names['tsrange'] = TSRANGE
+
+class TSTZRANGE(RangeOperators, sqltypes.TypeEngine):
+    "Represent the Postgresql TSTZRANGE type."
+    
+    __visit_name__ = 'TSTZRANGE'
+
+ischema_names['tstzrange'] = TSTZRANGE