bug-fix: adding new Kusto types also to the kql_converters (#75)

* bug-fix: adding new Kusto types also to the kql_converters

* supporting additional column types in PyKusto

* Assert that type registrars cover all types + some style improvements

* Deduplicate list of number types, support decimal type, some style improvements

* Added link

Co-authored-by: Netanel Zakay <nezakay@microsoft.com>
Co-authored-by: Yonatan Most <>

.idea/dictionaries/pykusto.xml

@@ -91,6 +91,7 @@
       <w>totimespan</w>
       <w>toupper</w>
       <w>treepath</w>
+      <w>uint</w>
       <w>urlquery</w>
       <w>varianceif</w>
       <w>variancep</w>

pykusto/expressions.py

@@ -4,7 +4,7 @@ from typing import Union
 
 from pykusto.keywords import KUSTO_KEYWORDS
 from pykusto.kql_converters import KQL
-from pykusto.type_utils import plain_expression, aggregation_expression, PythonTypes, kql_converter, KustoType, typed_column, TypeRegistrar, get_base_types
+from pykusto.type_utils import plain_expression, aggregation_expression, PythonTypes, kql_converter, KustoType, typed_column, TypeRegistrar, get_base_types, NUMBER_TYPES
 
 ExpressionType = Union[PythonTypes, 'BaseExpression']
 StringType = Union[str, 'StringExpression']
@@ -201,7 +201,7 @@ class BooleanExpression(BaseExpression):
         return BooleanExpression(KQL(f'not({self.kql})'))
 
 
-@plain_expression(KustoType.INT, KustoType.LONG, KustoType.REAL)
+@plain_expression(*NUMBER_TYPES)
 class NumberExpression(BaseExpression):
     @staticmethod
     def binary_op(left: NumberType, operator: str, right: NumberType) -> 'NumberExpression':
@@ -709,7 +709,7 @@ class BooleanAggregationExpression(AggregationExpression, BooleanExpression):
     pass
 
 
-@aggregation_expression(KustoType.INT, KustoType.LONG, KustoType.REAL)
+@aggregation_expression(*NUMBER_TYPES)
 class NumberAggregationExpression(AggregationExpression, NumberExpression):
     pass
 
@@ -807,7 +807,7 @@ class BaseColumn(BaseExpression):
         return self._name
 
 
-@typed_column(KustoType.INT, KustoType.LONG, KustoType.REAL)
+@typed_column(*NUMBER_TYPES)
 class NumberColumn(BaseColumn, NumberExpression):
     pass
 
@@ -850,9 +850,10 @@ class SubtractableColumn(NumberColumn, DatetimeColumn, TimespanColumn):
     def __sub__(self, other: Union['NumberType', 'DatetimeType', 'TimespanType']) -> Union['NumberExpression', 'TimespanExpression', 'AnyExpression']:
         # noinspection PyTypeChecker
         base_types = get_base_types(other)
-        possible_types = base_types & {KustoType.DATETIME, KustoType.TIMESPAN, KustoType.INT, KustoType.LONG, KustoType.REAL}
+        possible_types = base_types & ({KustoType.DATETIME, KustoType.TIMESPAN} | NUMBER_TYPES)
+
         assert len(possible_types) > 0, "Invalid type subtracted"
-        if possible_types == {KustoType.INT, KustoType.LONG, KustoType.REAL}:
+        if possible_types == NUMBER_TYPES:
             return_type = KustoType.INT
         elif len(possible_types) > 1:
             return_type = None
@@ -909,3 +910,8 @@ def to_kql(obj: ExpressionType) -> KQL:
 def expression_to_type(expression: ExpressionType, type_registrar: TypeRegistrar, fallback_type: Any) -> Any:
     types = set(type_registrar.registry[base_type] for base_type in plain_expression.get_base_types(expression))
     return next(iter(types)) if len(types) == 1 else fallback_type
+
+
+typed_column.assert_all_types_covered()
+plain_expression.assert_all_types_covered()
+aggregation_expression.assert_all_types_covered()

pykusto/kql_converters.py

@@ -4,7 +4,7 @@ from itertools import chain
 from numbers import Number
 from typing import NewType, Union, Mapping, List, Tuple
 
-from pykusto.type_utils import kql_converter, KustoType
+from pykusto.type_utils import kql_converter, KustoType, NUMBER_TYPES
 
 KQL = NewType('KQL', str)
 
@@ -73,6 +73,9 @@ def str_to_kql(s: str) -> KQL:
     return KQL(f'"{s}"')
 
 
-@kql_converter(KustoType.INT, KustoType.LONG, KustoType.REAL)
+@kql_converter(*NUMBER_TYPES)
 def number_to_kql(n: Number) -> KQL:
     return KQL(str(n))
+
+
+kql_converter.assert_all_types_covered()

pykusto/type_utils.py

@@ -1,35 +1,39 @@
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Union, Mapping, Type, Dict, Callable, Tuple, List, Set
+from typing import Union, Mapping, Type, Dict, Callable, Tuple, List, Set, FrozenSet
 
-# TODO: Unhandled data types: guid, decimal
 PythonTypes = Union[str, int, float, bool, datetime, Mapping, List, Tuple, timedelta]
 
 
 class KustoType(Enum):
+    """
+    https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/scalar-data-types/
+    """
     BOOL = ('bool', 'I8', 'System.SByte', bool)
     DATETIME = ('datetime', 'DateTime', 'System.DateTime', datetime)
-    # noinspection PyTypeChecker
-    DECIMAL = ('decimal', 'Decimal', 'System.Data.SqlTypes.SqlDecimal', None)  # TODO
     ARRAY = ('dynamic', 'Dynamic', 'System.Object', List, Tuple)
     MAPPING = ('dynamic', 'Dynamic', 'System.Object', Mapping)
-    # noinspection PyTypeChecker
-    GUID = ('guid', 'UniqueId', 'System.Guid', None)  # TODO
     INT = ('int', 'I32', 'System.Int32', int)
     LONG = ('long', 'I64', 'System.Int64', int)
     REAL = ('real', 'R64', 'System.Double', float)
     STRING = ('string', 'StringBuffer', 'System.String', str)
     TIMESPAN = ('timespan', 'TimeSpan', 'System.TimeSpan', timedelta)
-    NULL = ('null', 'null', 'null', type(None))
-    # the following types are not supported in Kusto anymore
-    UINT8 = ('uint8', 'UI8', 'System.Byte', int)
+    DECIMAL = ('decimal', 'Decimal', 'System.Data.SqlTypes.SqlDecimal', int)
+    GUID = ('guid', 'UniqueId', 'System.Guid')  # Not supported by Kusto yet
+
+    # Deprecated types, kept here for back compatibility
+    # https://docs.microsoft.com/en-us/azure/data-explorer/kusto/query/scalar-data-types/unsupported-data-types
+    FLOAT = ('float', 'R32', 'System.Single', float)
+    INT16 = ('int16', 'I16', 'System.Int16', int)
+    UINT16 = ('uint16', 'UI16', 'System.UInt16', int)
     UINT32 = ('uint32', 'UI32', 'System.UInt32', int)
-    UINT64 = ('UI64', 'UI64', 'System.UInt64', int)
+    UINT64 = ('uint64', 'UI64', 'System.UInt64', int)
+    UINT8 = ('uint8', 'UI8', 'System.Byte', int)
 
     primary_name: str
     internal_name: str
     dot_net_name: str
-    python_types: Tuple[PythonTypes]
+    python_types: Tuple[PythonTypes, ...]
 
     def __init__(self, primary_name: str, internal_name: str, dot_net_name: str, *python_types: PythonTypes) -> None:
         self.primary_name = primary_name
@@ -52,6 +56,9 @@ class KustoType(Enum):
 
 INTERNAL_NAME_TO_TYPE: Dict[str, KustoType] = {t.internal_name: t for t in KustoType}
 DOT_NAME_TO_TYPE: Dict[str, KustoType] = {t.dot_net_name: t for t in KustoType}
+NUMBER_TYPES: FrozenSet[KustoType] = frozenset([
+    KustoType.INT, KustoType.LONG, KustoType.REAL, KustoType.DECIMAL, KustoType.FLOAT, KustoType.INT16, KustoType.UINT16, KustoType.UINT32, KustoType.UINT64, KustoType.UINT8
+])
 
 
 class TypeRegistrar:
@@ -131,6 +138,10 @@ class TypeRegistrar:
         assert len(base_types) > 0, f"get_base_types called for unsupported type: {type(obj).__name__}"
         return base_types
 
+    def assert_all_types_covered(self) -> None:
+        missing = set(t for t in KustoType if len(t.python_types) > 0) - set(self.registry.keys())
+        assert len(missing) == 0, [t.name for t in missing]
+
 
 kql_converter = TypeRegistrar("KQL Converter")
 typed_column = TypeRegistrar("Column")

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='pykusto',
-    version='0.0.18',
+    version='0.0.19',
     packages=find_packages(exclude=['test']),
     url='https://github.com/Azure/pykusto',
     license='MIT License',