Add regexp operators for PostgreSQL string types.pr/281

3 files changed, 173 insertions, 4 deletions

diff --git a/doc/build/dialects/postgresql.rst b/doc/build/dialects/postgresql.rst
index b4c90643d..9ba76df2b 100644
--- a/doc/build/dialects/postgresql.rst
+++ b/doc/build/dialects/postgresql.rst
@@ -44,6 +44,8 @@ construction arguments, are as follows:
 .. autoclass:: BYTEA
     :members: __init__
 
+.. autoclass:: CHAR
+
 .. autoclass:: CIDR
 
 
@@ -85,12 +87,21 @@ construction arguments, are as follows:
 .. autoclass:: REAL
     :members: __init__
 
+.. autoclass:: TEXT
+
 .. autoclass:: TSVECTOR
     :members: __init__
 
 .. autoclass:: UUID
     :members: __init__
 
+.. autoclass:: VARCHAR
+
+The string types get additional functionality from the following mixin:
+
+.. autoclass:: sqlalchemy.dialects.postgresql.base._PGStringOps
+    :members: Comparator
+
 
 Range Types
 ~~~~~~~~~~~
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index 1bc4409f2..69472e8fa 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -619,7 +619,7 @@ import datetime as dt
 
 from ... import sql, schema, exc, util
 from ...engine import default, reflection
-from ...sql import compiler, expression
+from ...sql import compiler, default_comparator, expression
 from ... import types as sqltypes
 
 try:
@@ -627,8 +627,7 @@ try:
 except ImportError:
     _python_UUID = None
 
-from sqlalchemy.types import INTEGER, BIGINT, SMALLINT, VARCHAR, \
-    CHAR, TEXT, FLOAT, NUMERIC, \
+from sqlalchemy.types import INTEGER, BIGINT, SMALLINT, FLOAT, NUMERIC, \
     DATE, BOOLEAN, REAL
 
 RESERVED_WORDS = set(
@@ -654,6 +653,67 @@ _FLOAT_TYPES = (700, 701, 1021, 1022)
 _INT_TYPES = (20, 21, 23, 26, 1005, 1007, 1016)
 
 
+def regexp_op(a, b):
+    raise NotImplementedError()
+
+
+def iregexp_op(a, b):
+    raise NotImplementedError()
+
+
+def notregexp_op(a, b):
+    raise NotImplementedError()
+
+
+def notiregexp_op(a, b):
+    raise NotImplementedError()
+
+
+class _PGStringOps(object):
+    """This mixin provides functionality for the POSIX regular expresssion
+    operators listed in Table 9-12 of the postgres documentation. It is
+    used by all the string types provided in the postgres dialect.
+    """
+    class Comparator(sqltypes.String.Comparator):
+        """Define comparison operations for string types."""
+        def regexp(self, other):
+            """Boolean expression. Returns true if the column text matches
+            the given regular expression (case sensitive).
+            """
+            return default_comparator._boolean_compare(
+                self.expr, regexp_op, other, negate=notregexp_op)
+
+        def iregexp(self, other):
+            """Boolean expression. Returns true if the column text matches
+            the given regular expression (case insensitive).
+            """
+            return default_comparator._boolean_compare(
+                self.expr, iregexp_op, other, negate=notiregexp_op)
+
+    comparator_factory = Comparator
+
+
+class CHAR(_PGStringOps, sqltypes.CHAR):
+    """Implement the CHAR type, adding in Postgresql-specific string
+    operators.
+
+    """
+
+
+class VARCHAR(_PGStringOps, sqltypes.VARCHAR):
+    """Implement the VARCHAR type, adding in Postgresql-specific string
+    operators.
+
+    """
+
+
+class TEXT(_PGStringOps, sqltypes.TEXT):
+    """Implement the TEXT type, adding in Postgresql-specific string
+    operators.
+
+    """
+
+
 class BYTEA(sqltypes.LargeBinary):
     __visit_name__ = 'BYTEA'
 
@@ -1089,6 +1149,24 @@ class PGCompiler(compiler.SQLCompiler):
             self.process(element.order_by, **kw)
         )
 
+    def visit_regexp_op_binary(self, binary, operator, **kw):
+        return self._binary_op(binary, "~", **kw)
+
+    def visit_iregexp_op_binary(self, binary, operator, **kw):
+        return self._binary_op(binary, "~*", **kw)
+
+    def visit_notregexp_op_binary(self, binary, operator, **kw):
+        return self._binary_op(binary, "!~", **kw)
+
+    def visit_notiregexp_op_binary(self, binary, operator, **kw):
+        return self._binary_op(binary, "!~*", **kw)
+
+    def _binary_op(self, binary, opstring, **kw):
+        return "%s %s %s" % (
+                        self.process(binary.left, **kw),
+                        opstring,
+                        self.process(binary.right, **kw))
+
     def visit_match_op_binary(self, binary, operator, **kw):
         if "postgresql_regconfig" in binary.modifiers:
             regconfig = self.render_literal_value(
diff --git a/test/dialect/postgresql/test_types.py b/test/dialect/postgresql/test_types.py
index 6bcc4cf9a..85d4b9d71 100644
--- a/test/dialect/postgresql/test_types.py
+++ b/test/dialect/postgresql/test_types.py
@@ -14,7 +14,7 @@ import sqlalchemy as sa
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects.postgresql import HSTORE, hstore, array, \
     INT4RANGE, INT8RANGE, NUMRANGE, DATERANGE, TSRANGE, TSTZRANGE, \
-    JSON, JSONB
+    JSON, JSONB, CHAR, VARCHAR, TEXT
 import decimal
 from sqlalchemy import util
 from sqlalchemy.testing.util import round_decimal
@@ -2402,6 +2402,86 @@ class DateTimeTZRangeTests(_RangeTypeMixin, fixtures.TablesTest):
         return self.extras().DateTimeTZRange(*self.tstzs())
 
 
+class RegexpTest(AssertsCompiledSQL, fixtures.TestBase):
+    __dialect__ = 'postgresql'
+
+    def setup(self):
+        metadata = MetaData()
+        self.test_table = Table('test_table', metadata,
+                                Column('id', Integer, primary_key=True),
+                                Column('char_col', CHAR),
+                                Column('varchar_col', VARCHAR),
+                                Column('text_col', TEXT))
+
+    def _test_where(self, whereclause, expected):
+        stmt = select([self.test_table]).where(whereclause)
+        self.assert_compile(
+            stmt,
+            "SELECT test_table.id, test_table.char_col, test_table."
+            "varchar_col, test_table.text_col FROM test_table WHERE %s" %
+            expected)
+
+    def test_char_regexp(self):
+        self._test_where(
+            self.test_table.c.char_col.regexp('.*'),
+            "(test_table.char_col ~ %(char_col_1)s)")
+
+    def test_varchar_regexp(self):
+        self._test_where(
+            self.test_table.c.varchar_col.regexp('.*'),
+            "(test_table.varchar_col ~ %(varchar_col_1)s)")
+
+    def test_text_regexp(self):
+        self._test_where(
+            self.test_table.c.text_col.regexp('.*'),
+            "(test_table.text_col ~ %(text_col_1)s)")
+
+    def test_char_iregexp(self):
+        self._test_where(
+            self.test_table.c.char_col.iregexp('.*'),
+            "(test_table.char_col ~* %(char_col_1)s)")
+
+    def test_varchar_iregexp(self):
+        self._test_where(
+            self.test_table.c.varchar_col.iregexp('.*'),
+            "(test_table.varchar_col ~* %(varchar_col_1)s)")
+
+    def test_text_iregexp(self):
+        self._test_where(
+            self.test_table.c.text_col.iregexp('.*'),
+            "(test_table.text_col ~* %(text_col_1)s)")
+
+    def test_char_not_regexp(self):
+        self._test_where(
+            ~self.test_table.c.char_col.regexp('.*'),
+            "(test_table.char_col !~ %(char_col_1)s)")
+
+    def test_varchar_not_regexp(self):
+        self._test_where(
+            ~self.test_table.c.varchar_col.regexp('.*'),
+            "(test_table.varchar_col !~ %(varchar_col_1)s)")
+
+    def test_text_not_regexp(self):
+        self._test_where(
+            ~self.test_table.c.text_col.regexp('.*'),
+            "(test_table.text_col !~ %(text_col_1)s)")
+
+    def test_char_not_iregexp(self):
+        self._test_where(
+            ~self.test_table.c.char_col.iregexp('.*'),
+            "(test_table.char_col !~* %(char_col_1)s)")
+
+    def test_varchar_not_iregexp(self):
+        self._test_where(
+            ~self.test_table.c.varchar_col.iregexp('.*'),
+            "(test_table.varchar_col !~* %(varchar_col_1)s)")
+
+    def test_text_not_iregexp(self):
+        self._test_where(
+            ~self.test_table.c.text_col.iregexp('.*'),
+            "(test_table.text_col !~* %(text_col_1)s)")
+
+
 class JSONTest(AssertsCompiledSQL, fixtures.TestBase):
     __dialect__ = 'postgresql'