Introduce typings for ConditionTranslator

setup.py

@@ -44,6 +44,7 @@
     'tenacity',
     'toucan_data_sdk',
     'urllib3==1.24.3',
+    'typing-extensions; python_version < "3.8"',
 ]
 
 classifiers = [

tests/mongo/test_mongo_translator.py

@@ -14,11 +14,11 @@ def test_translate_condition_unit_to_mongo_match():
     c = {'column': 'country', 'operator': 'eq', 'value': 'France'}
     assert MongoConditionTranslator.translate(c) == {'country': {'$eq': 'France'}}
     # raise when needed
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         MongoConditionTranslator.translate({'column': 'population', 'operator': 'eq'})
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         MongoConditionTranslator.translate({'column': 'population', 'value': 42})
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         MongoConditionTranslator.translate({'operator': 'eq', 'value': 42})
     with pytest.raises(ValueError):
         MongoConditionTranslator.translate(
@@ -130,10 +130,10 @@ def test_MongoConditionTranslator_translate_with_jinja():
 def test_MongoConditionTranslator_operators():
     assert MongoConditionTranslator.EQUAL('col', 'val') == {'col': {'$eq': 'val'}}
     assert MongoConditionTranslator.NOT_EQUAL('col', 'val') == {'col': {'$ne': 'val'}}
-    assert MongoConditionTranslator.GREATER_THAN('col', 'val') == {'col': {'$gt': 'val'}}
-    assert MongoConditionTranslator.GREATER_THAN_EQUAL('col', 'val') == {'col': {'$gte': 'val'}}
-    assert MongoConditionTranslator.LOWER_THAN('col', 'val') == {'col': {'$lt': 'val'}}
-    assert MongoConditionTranslator.LOWER_THAN_EQUAL('col', 'val') == {'col': {'$lte': 'val'}}
+    assert MongoConditionTranslator.GREATER_THAN('col', 3) == {'col': {'$gt': 3}}
+    assert MongoConditionTranslator.GREATER_THAN_EQUAL('col', 3) == {'col': {'$gte': 3}}
+    assert MongoConditionTranslator.LOWER_THAN('col', 3) == {'col': {'$lt': 3}}
+    assert MongoConditionTranslator.LOWER_THAN_EQUAL('col', 3) == {'col': {'$lte': 3}}
     assert MongoConditionTranslator.IN('col', ['val']) == {'col': {'$in': ['val']}}
     assert MongoConditionTranslator.NOT_IN('col', ['val']) == {'col': {'$nin': ['val']}}
     assert MongoConditionTranslator.IS_NULL('col') == {'col': {'$exists': False}}

tests/test_pandas_translator.py

@@ -14,11 +14,11 @@ def test_translate_condition_unit():
     c = {'column': 'country', 'operator': 'eq', 'value': 'France'}
     assert PandasConditionTranslator.translate(c) == "`country` == 'France'"
     # error cases
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         PandasConditionTranslator.translate({'column': 'population', 'operator': 'eq'})
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         PandasConditionTranslator.translate({'column': 'population', 'value': 42})
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         PandasConditionTranslator.translate({'operator': 'eq', 'value': 42})
     with pytest.raises(ValueError):
         PandasConditionTranslator.translate(

toucan_connectors/condition_translator.py

@@ -1,8 +1,24 @@
+import sys
 from abc import ABC, abstractmethod
 from enum import Enum
+from typing import Any, List, TypeVar, Union
 
+from pydantic import BaseModel
 
-class ConditionOperator(str, Enum):
+if sys.version_info < (3, 8):
+    from typing_extensions import Literal
+else:
+    from typing import Literal
+
+Number = Union[int, float]
+
+ClauseType = TypeVar('ClauseType')
+
+
+LogicalOperator = Union[Literal['and'], Literal['or']]
+
+
+class ConditionOperator(Enum):
     EQUAL = 'eq'
     NOT_EQUAL = 'ne'
     LOWER_THAN = 'lt'
@@ -17,6 +33,12 @@ class ConditionOperator(str, Enum):
     IS_NOT_NULL = 'notnull'
 
 
+class ConditionModel(BaseModel):
+    column: str
+    operator: ConditionOperator
+    value: Any = ...
+
+
 class ConditionTranslator(ABC):
     """
     Class with utilities methods to translate data condition from a
@@ -26,7 +48,7 @@ class ConditionTranslator(ABC):
     """
 
     @classmethod
-    def translate(cls, condition: dict, **kwargs):
+    def translate(cls, condition: dict):
         """
         Convert a condition into a format relevant for a type of connector.
 
@@ -52,100 +74,116 @@ def translate(cls, condition: dict, **kwargs):
         if 'or' in condition:
             if isinstance(condition['or'], list):
                 return cls.join_clauses(
-                    [cls.translate(condition, **kwargs) for condition in condition['or']], 'or'
+                    [cls.translate(condition) for condition in condition['or']], 'or'
                 )
             else:
                 raise ValueError("'or' value must be an array")
         elif 'and' in condition:
             if isinstance(condition['and'], list):
                 return cls.join_clauses(
-                    [cls.translate(condition, **kwargs) for condition in condition['and']], 'and'
+                    [cls.translate(condition) for condition in condition['and']], 'and'
                 )
             else:
                 raise ValueError("'and' value must be an array")
         else:
-            return cls.generate_clause(**condition, **kwargs)
+            condition_m = ConditionModel(**condition)
+            clause_generator_for_operator = getattr(cls, condition_m.operator.name)
+
+            if isinstance(condition_m.value, str):
+                condition_m.value = cls.get_value_str_ref(condition_m.value)
+
+            return clause_generator_for_operator(
+                cls.get_column_ref(condition_m.column), condition_m.value
+            )
 
     @classmethod
     @abstractmethod
-    def join_clauses(cls, clauses: list, logical_operator: str):
+    def join_clauses(cls, clauses: List[ClauseType], logical_operator: LogicalOperator):
         """
         Join multiple clauses with `and` or `or`.
         """
         raise NotImplementedError
 
     @classmethod
-    def generate_clause(
-        cls, column: str, operator: str, value, enclosing_field_char='', enclosing_value_char=''
-    ):
-        condition_operator = ConditionOperator(operator)
-        clause_generator_for_operator = getattr(cls, condition_operator.name)
+    def get_column_ref(cls, column_name: str) -> str:
+        """How to refer to column in the clause"""
+        return column_name
 
-        if isinstance(value, str):
-            value = f'{enclosing_value_char}{value}{enclosing_value_char}'
-
-        return clause_generator_for_operator(
-            f'{enclosing_field_char}{column}{enclosing_field_char}', value
-        )
-
-    # Operators
+    @classmethod
+    def get_value_str_ref(cls, value: str) -> str:
+        """How to refer to value strings in the clause"""
+        return value
 
+    # Operators to implement
+    # `column` and `value` are ref strings from `get_column_ref` and `get_value_str_ref`
     @classmethod
     @abstractmethod
-    def EQUAL(cls, column, value):
+    def EQUAL(cls, column: str, value: str) -> ClauseType:
+        """`column` values equal to `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def NOT_EQUAL(cls, column, value):
+    def NOT_EQUAL(cls, column: str, value: str) -> ClauseType:
+        """`column` values not equal to `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def LOWER_THAN(cls, column, value):
+    def LOWER_THAN(cls, column: str, value: Number) -> ClauseType:
+        """`column` values lower than `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def LOWER_THAN_EQUAL(cls, column, value):
+    def LOWER_THAN_EQUAL(cls, column: str, value: Number) -> ClauseType:
+        """`column` values lower than or equal to `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def GREATER_THAN(cls, column, value):
+    def GREATER_THAN(cls, column: str, value: Number) -> ClauseType:
+        """`column` values greater than `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def GREATER_THAN_EQUAL(cls, column, value):
+    def GREATER_THAN_EQUAL(cls, column: str, value: Number) -> ClauseType:
+        """`column` values greater than or equal to `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def IN(cls, column, values):
+    def IN(cls, column: str, values: List[str]) -> ClauseType:
+        """`column` values in `values`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def NOT_IN(cls, column, values):
+    def NOT_IN(cls, column: str, values: List[str]) -> ClauseType:
+        """`column` values not in `values`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def MATCHES(cls, column, value):
+    def MATCHES(cls, column: str, value: str) -> ClauseType:
+        """`column` values match the regex `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def NOT_MATCHES(cls, column, value):
+    def NOT_MATCHES(cls, column, value) -> ClauseType:
+        """`column` values don't match the regex `value`"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def IS_NULL(cls, column, value=None):
+    def IS_NULL(cls, column: str, value=None) -> ClauseType:
+        """`column` values are null"""
         raise NotImplementedError
 
     @classmethod
     @abstractmethod
-    def IS_NOT_NULL(cls, column, value=None):
+    def IS_NOT_NULL(cls, column: str, value=None) -> ClauseType:
+        """`column` values are not null"""
         raise NotImplementedError

toucan_connectors/mongo/mongo_translator.py

@@ -1,3 +1,5 @@
+from typing import Dict, List
+
 from toucan_connectors.condition_translator import ConditionTranslator
 
 
@@ -7,53 +9,53 @@ class MongoConditionTranslator(ConditionTranslator):
     """
 
     @classmethod
-    def join_clauses(cls, clauses: list, logical_operator: str):
+    def join_clauses(cls, clauses: List[dict], logical_operator: str):
         return {f'${logical_operator}': clauses}
 
     @classmethod
-    def EQUAL(cls, column, value):
+    def EQUAL(cls, column, value) -> Dict[str, dict]:
         return {column: {'$eq': value}}
 
     @classmethod
-    def NOT_EQUAL(cls, column, value):
+    def NOT_EQUAL(cls, column, value) -> Dict[str, dict]:
         return {column: {'$ne': value}}
 
     @classmethod
-    def LOWER_THAN(cls, column, value):
+    def LOWER_THAN(cls, column, value) -> Dict[str, dict]:
         return {column: {'$lt': value}}
 
     @classmethod
-    def LOWER_THAN_EQUAL(cls, column, value):
+    def LOWER_THAN_EQUAL(cls, column, value) -> Dict[str, dict]:
         return {column: {'$lte': value}}
 
     @classmethod
-    def GREATER_THAN(cls, column, value):
+    def GREATER_THAN(cls, column, value) -> Dict[str, dict]:
         return {column: {'$gt': value}}
 
     @classmethod
-    def GREATER_THAN_EQUAL(cls, column, value):
+    def GREATER_THAN_EQUAL(cls, column, value) -> Dict[str, dict]:
         return {column: {'$gte': value}}
 
     @classmethod
-    def IN(cls, column, values):
+    def IN(cls, column, values) -> Dict[str, dict]:
         return {column: {'$in': values}}
 
     @classmethod
-    def NOT_IN(cls, column, values):
+    def NOT_IN(cls, column, values) -> Dict[str, dict]:
         return {column: {'$nin': values}}
 
     @classmethod
-    def MATCHES(cls, column, value):
+    def MATCHES(cls, column, value) -> Dict[str, dict]:
         return {column: {'$regex': value}}
 
     @classmethod
-    def NOT_MATCHES(cls, column, value):
+    def NOT_MATCHES(cls, column, value) -> Dict[str, dict]:
         return {column: {'$not': {'$regex': value}}}
 
     @classmethod
-    def IS_NULL(cls, column, value=None):
+    def IS_NULL(cls, column, value=None) -> Dict[str, dict]:
         return {column: {'$exists': False}}
 
     @classmethod
-    def IS_NOT_NULL(cls, column, value=None):
+    def IS_NOT_NULL(cls, column, value=None) -> Dict[str, dict]:
         return {column: {'$exists': True}}

toucan_connectors/pandas_translator.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from toucan_connectors.condition_translator import ConditionTranslator
 
 
@@ -10,45 +12,46 @@ class PandasConditionTranslator(ConditionTranslator):
     """
 
     @classmethod
-    def translate(cls, conditions: dict, enclosing_field_char='`', enclosing_value_char="'") -> str:
-        return super().translate(
-            conditions,
-            enclosing_field_char=enclosing_field_char,
-            enclosing_value_char=enclosing_value_char,
-        )
+    def get_column_ref(cls, column: str) -> str:
+        """To refer column names (even with spaces or operators), we surround them in backticks"""
+        return f'`{column}`'
+
+    @classmethod
+    def get_value_str_ref(cls, value: str) -> str:
+        return f"'{value}'"
 
     @classmethod
-    def join_clauses(cls, clauses: list, logical_operator: str):
+    def join_clauses(cls, clauses: List[str], logical_operator: str) -> str:
         return '(' + f' {logical_operator} '.join(clauses) + ')'
 
     @classmethod
-    def EQUAL(cls, column, value):
+    def EQUAL(cls, column, value) -> str:
         return f'{column} == {value}'
 
     @classmethod
-    def NOT_EQUAL(cls, column, value):
+    def NOT_EQUAL(cls, column, value) -> str:
         return f'{column} != {value}'
 
     @classmethod
-    def LOWER_THAN(cls, column, value):
+    def LOWER_THAN(cls, column, value) -> str:
         return f'{column} < {value}'
 
     @classmethod
-    def LOWER_THAN_EQUAL(cls, column, value):
+    def LOWER_THAN_EQUAL(cls, column, value) -> str:
         return f'{column} <= {value}'
 
     @classmethod
-    def GREATER_THAN(cls, column, value):
+    def GREATER_THAN(cls, column, value) -> str:
         return f'{column} > {value}'
 
     @classmethod
-    def GREATER_THAN_EQUAL(cls, column, value):
+    def GREATER_THAN_EQUAL(cls, column, value) -> str:
         return f'{column} >= {value}'
 
     @classmethod
-    def IN(cls, column, value):
+    def IN(cls, column, value) -> str:
         return f'{column} in {value}'
 
     @classmethod
-    def NOT_IN(cls, column, value):
+    def NOT_IN(cls, column, value) -> str:
         return f'{column} not in {value}'