Repository: Azure/pykusto

Issue:
Return correct types in operators that accept aggregation expressions

Discussion:

Patch:
diff --git a/pykusto/client.py b/pykusto/client.py
index 2bc47e4..23b7266 100644
--- a/pykusto/client.py
+++ b/pykusto/client.py
@@ -1,5 +1,4 @@
 from typing import Union, List, Tuple
-
 # noinspection PyProtectedMember
 from urllib.parse import urlparse
 
@@ -7,7 +6,7 @@ from urllib.parse import urlparse
 from azure.kusto.data._response import KustoResponseDataSet
 from azure.kusto.data.request import KustoClient, KustoConnectionStringBuilder, ClientRequestProperties
 
-from pykusto.utils import KQL
+from pykusto.kql_converters import KQL
 
 
 class PyKustoClient:
diff --git a/pykusto/expressions.py b/pykusto/expressions.py
index f1e9db8..a02183b 100644
--- a/pykusto/expressions.py
+++ b/pykusto/expressions.py
@@ -1,9 +1,10 @@
 from datetime import datetime, timedelta
-from typing import Any, List, Tuple, Mapping, Optional
+from numbers import Number
+from typing import Any, List, Tuple, Mapping, Optional, Type
 from typing import Union
 
-from pykusto.utils import KQL
-from pykusto.utils import KustoTypes, to_kql
+from pykusto.kql_converters import KQL
+from pykusto.type_utils import plain_expression, aggregation_expression, KustoTypes, kql_converter
 
 ExpressionType = Union[KustoTypes, 'BaseExpression']
 StringType = Union[str, 'StringExpression']
@@ -13,14 +14,21 @@ ArrayType = Union[List, Tuple, 'ArrayExpression']
 MappingType = Union[Mapping, 'MappingExpression']
 DatetimeType = Union[datetime, 'DatetimeExpression']
 TimespanType = Union[timedelta, 'TimespanExpression']
-AggregationType = Union['AggregationExpression']
-DynamicType = Union['DynamicExpression']
+DynamicType = Union[ArrayType, MappingType]
 OrderType = Union[DatetimeType, TimespanType, NumberType, StringType]
 
 
 # All classes in the same file to prevent circular dependencies
 
 def _subexpr_to_kql(obj: ExpressionType) -> KQL:
+    """
+    Convert the given expression to KQL, enclosing it in parentheses if it is a compound expression. This guarantees
+    correct evaluation order. When parentheses are not needed, for example when the expressions is used as an argument
+    to a function, use `to_kql` instead.
+
+    :param obj: Expression to convert to KQL
+    :return: KQL that represents the given expression
+    """
     if isinstance(obj, BaseExpression):
         return obj.as_subexpression()
     return to_kql(obj)
@@ -29,11 +37,20 @@ def _subexpr_to_kql(obj: ExpressionType) -> KQL:
 class BaseExpression:
     kql: KQL
 
-    def __init__(self, kql: KQL) -> None:
+    def __new__(cls, *args, **kwargs):
+        if cls is 'BaseExpression':
+            raise TypeError("BaseExpression is abstract")
+        return object.__new__(cls)
+
+    def __init__(self, kql: Union[KQL, 'BaseExpression']) -> None:
+        if isinstance(kql, BaseExpression):
+            self.kql = kql.kql
+        elif not isinstance(kql, str):
+            raise ValueError("Either expression or KQL required")
         self.kql = kql
 
     def __repr__(self) -> str:
-        return self.kql
+        return str(self.kql)
 
     def as_subexpression(self) -> KQL:
         return KQL('({})'.format(self.kql))
@@ -58,10 +75,15 @@ class BaseExpression:
         return BooleanExpression(KQL('{} has \"{}\"'.format(self.kql, exp)))
 
     @staticmethod
-    def binary_op(left: ExpressionType, operator: str, right: ExpressionType) -> KQL:
-        return KQL('{}{}{}'.format(
+    def base_binary_op(
+            left: ExpressionType, operator: str, right: ExpressionType, result_type: Type[KustoTypes]
+    ) -> 'BaseExpression':
+        registrar = plain_expression
+        if isinstance(left, AggregationExpression) or isinstance(right, AggregationExpression):
+            registrar = aggregation_expression
+        return registrar.for_type(result_type)(KQL('{}{}{}'.format(
             _subexpr_to_kql(left), operator, _subexpr_to_kql(right))
-        )
+        ))
 
     def __eq__(self, other: ExpressionType) -> 'BooleanExpression':
         return BooleanExpression.binary_op(self, ' == ', other)
@@ -100,10 +122,12 @@ class BaseExpression:
         raise ValueError("Only arrays can be assigned to multiple columns")
 
 
+@plain_expression(bool)
 class BooleanExpression(BaseExpression):
     @staticmethod
     def binary_op(left: ExpressionType, operator: str, right: ExpressionType) -> 'BooleanExpression':
-        return BooleanExpression(BaseExpression.binary_op(left, operator, right))
+        # noinspection PyTypeChecker
+        return BaseExpression.base_binary_op(left, operator, right, bool)
 
     def __and__(self, other: BooleanType) -> 'BooleanExpression':
         return BooleanExpression.binary_op(self, ' and ', other)
@@ -115,10 +139,12 @@ class BooleanExpression(BaseExpression):
         return BooleanExpression(KQL('not({})'.format(self.kql)))
 
 
+@plain_expression(Number)
 class NumberExpression(BaseExpression):
     @staticmethod
     def binary_op(left: NumberType, operator: str, right: NumberType) -> 'NumberExpression':
-        return NumberExpression(BaseExpression.binary_op(left, operator, right))
+        # noinspection PyTypeChecker
+        return BaseExpression.base_binary_op(left, operator, right, Number)
 
     def __lt__(self, other: NumberType) -> BooleanExpression:
         return BooleanExpression.binary_op(self, ' < ', other)
@@ -168,15 +194,15 @@ class NumberExpression(BaseExpression):
         return NumberExpression(KQL('floor({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
 
     def bin(self, round_to: NumberType) -> 'BaseExpression':
-        return BaseExpression(KQL('bin({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
+        return NumberExpression(KQL('bin({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
 
     def bin_at(self, round_to: NumberType, fixed_point: NumberType) -> 'BaseExpression':
-        return BaseExpression(KQL('bin_at({}, {}, {})'.format(self.kql,
-                                                              _subexpr_to_kql(round_to),
-                                                              _subexpr_to_kql(fixed_point))))
+        return NumberExpression(KQL('bin_at({}, {}, {})'.format(self.kql,
+                                                                _subexpr_to_kql(round_to),
+                                                                _subexpr_to_kql(fixed_point))))
 
     def bin_auto(self) -> 'BaseExpression':
-        return BaseExpression(KQL('bin_auto({})'.format(self.kql)))
+        return NumberExpression(KQL('bin_auto({})'.format(self.kql)))
 
     def ceiling(self) -> 'NumberExpression':
         return NumberExpression(KQL('ceiling({})'.format(self.kql)))
@@ -213,11 +239,17 @@ class NumberExpression(BaseExpression):
 
     def round(self, precision: NumberType = None) -> 'NumberExpression':
         return NumberExpression(KQL(
-            ('round({}, {})' if precision is None else 'round({}, {})').format(self, precision)
+            ('round({}, {})' if precision is None else 'round({}, {})').format(self.kql, to_kql(precision))
         ))
 
 
+@plain_expression(str)
 class StringExpression(BaseExpression):
+    @staticmethod
+    def binary_op(left: ExpressionType, operator: str, right: ExpressionType) -> 'StringExpression':
+        # noinspection PyTypeChecker
+        return BaseExpression.base_binary_op(left, operator, right, str)
+
     def __len__(self) -> NumberExpression:
         return self.string_size()
 
@@ -228,7 +260,7 @@ class StringExpression(BaseExpression):
         return BooleanExpression(KQL('isempty({})'.format(self.kql)))
 
     def __add__(self, other: StringType) -> 'StringExpression':
-        return StringExpression(BaseExpression.binary_op(self, ' + ', other))
+        return StringExpression.binary_op(self, ' + ', other)
 
     @staticmethod
     def concat(*strings: StringType) -> 'StringExpression':
@@ -238,8 +270,10 @@ class StringExpression(BaseExpression):
 
     def split(self, delimiter: StringType, requested_index: NumberType = None) -> 'ArrayExpression':
         if requested_index is None:
-            return ArrayExpression(KQL('split({}, {})'.format(_subexpr_to_kql(self.kql), _subexpr_to_kql(delimiter))))
-        return ArrayExpression(KQL('split({}, {}, {})'.format(_subexpr_to_kql(self.kql), _subexpr_to_kql(delimiter), requested_index)))
+            return ArrayExpression(KQL('split({}, {})'.format(to_kql(self.kql), to_kql(delimiter))))
+        return ArrayExpression(KQL('split({}, {}, {})'.format(
+            to_kql(self.kql), to_kql(delimiter), to_kql(requested_index)
+        )))
 
     def equals(self, other: StringType, case_sensitive: bool = False) -> BooleanExpression:
         return BooleanExpression.binary_op(self, ' == ' if case_sensitive else ' =~ ', other)
@@ -278,10 +312,12 @@ class StringExpression(BaseExpression):
         return BooleanExpression(KQL('isutf8({})'.format(self.kql)))
 
 
+@plain_expression(datetime)
 class DatetimeExpression(BaseExpression):
     @staticmethod
     def binary_op(left: ExpressionType, operator: str, right: ExpressionType) -> 'DatetimeExpression':
-        return DatetimeExpression(BaseExpression.binary_op(left, operator, right))
+        # noinspection PyTypeChecker
+        return BaseExpression.base_binary_op(left, operator, right, datetime)
 
     def __lt__(self, other: DatetimeType) -> BooleanExpression:
         return BooleanExpression.binary_op(self, ' < ', other)
@@ -316,15 +352,15 @@ class DatetimeExpression(BaseExpression):
         return DatetimeExpression(KQL('floor({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
 
     def bin(self, round_to: TimespanType) -> 'BaseExpression':
-        return BaseExpression(KQL('bin({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
+        return DatetimeExpression(KQL('bin({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
 
     def bin_at(self, round_to: TimespanType, fixed_point: DatetimeType) -> 'BaseExpression':
-        return BaseExpression(KQL('bin_at({}, {}, {})'.format(self.kql,
-                                                              _subexpr_to_kql(round_to),
-                                                              _subexpr_to_kql(fixed_point))))
+        return DatetimeExpression(KQL('bin_at({}, {}, {})'.format(
+            self.kql, _subexpr_to_kql(round_to), to_kql(fixed_point)
+        )))
 
     def bin_auto(self) -> 'BaseExpression':
-        return BaseExpression(KQL('bin_auto({})'.format(self.kql)))
+        return DatetimeExpression(KQL('bin_auto({})'.format(self.kql)))
 
     def endofday(self, offset: NumberType = None) -> 'DatetimeExpression':
         if offset is None:
@@ -364,33 +400,35 @@ class DatetimeExpression(BaseExpression):
         return NumberExpression(KQL('getyear({})'.format(self.kql)))
 
     def hourofday(self) -> NumberExpression:
-        return NumberExpression(KQL('hourofday({})'.format(self)))
+        return NumberExpression(KQL('hourofday({})'.format(self.kql)))
 
     def startofday(self, offset: NumberType = None) -> 'DatetimeExpression':
         return DatetimeExpression(KQL(
-            ('startofday({})' if offset is None else 'startofday({}, {})').format(self.kql, offset)
+            ('startofday({})' if offset is None else 'startofday({}, {})').format(self.kql, to_kql(offset))
         ))
 
     def startofmonth(self, offset: NumberType = None) -> 'DatetimeExpression':
         return DatetimeExpression(KQL(
-            ('startofmonth({})' if offset is None else 'startofmonth({}, {})').format(self.kql, offset)
+            ('startofmonth({})' if offset is None else 'startofmonth({}, {})').format(self.kql, to_kql(offset))
         ))
 
     def startofweek(self, offset: NumberType = None) -> 'DatetimeExpression':
         return DatetimeExpression(KQL(
-            ('startofweek({})' if offset is None else 'startofweek({}, {})').format(self.kql, offset)
+            ('startofweek({})' if offset is None else 'startofweek({}, {})').format(self.kql, to_kql(offset))
         ))
 
     def startofyear(self, offset: NumberType = None) -> 'DatetimeExpression':
         return DatetimeExpression(KQL(
-            ('startofyear({})' if offset is None else 'startofyear({}, {})').format(self.kql, offset)
+            ('startofyear({})' if offset is None else 'startofyear({}, {})').format(self.kql, to_kql(offset))
         ))
 
 
+@plain_expression(timedelta)
 class TimespanExpression(BaseExpression):
     @staticmethod
     def binary_op(left: ExpressionType, operator: str, right: ExpressionType) -> 'TimespanExpression':
-        return TimespanExpression(BaseExpression.binary_op(left, operator, right))
+        # noinspection PyTypeChecker
+        return BaseExpression.base_binary_op(left, operator, right, timedelta)
 
     def __add__(self, other: TimespanType) -> 'TimespanExpression':
         return TimespanExpression.binary_op(self, ' + ', other)
@@ -402,18 +440,18 @@ class TimespanExpression(BaseExpression):
         return DatetimeExpression(KQL('ago({})'.format(_subexpr_to_kql(self))))
 
     def bin(self, round_to: TimespanType) -> 'BaseExpression':
-        return BaseExpression(KQL('bin({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
+        return TimespanExpression(KQL('bin({}, {})'.format(self.kql, _subexpr_to_kql(round_to))))
 
     def bin_at(self, round_to: TimespanType, fixed_point: TimespanType) -> 'BaseExpression':
-        return BaseExpression(KQL('bin_at({}, {}, {})'.format(self.kql,
-                                                              _subexpr_to_kql(round_to),
-                                                              _subexpr_to_kql(fixed_point))))
+        return TimespanExpression(KQL('bin_at({}, {}, {})'.format(
+            self.kql, to_kql(round_to), to_kql(fixed_point)
+        )))
 
     def bin_auto(self) -> 'BaseExpression':
-        return BaseExpression(KQL('bin_auto({})'.format(self.kql)))
+        return TimespanExpression(KQL('bin_auto({})'.format(self.kql)))
 
     def format_timespan(self, format_string: StringType) -> StringExpression:
-        return StringExpression(KQL('format_timespan({}, {})'.format(self.kql, _subexpr_to_kql(format_string))))
+        return StringExpression(KQL('format_timespan({}, {})'.format(self.kql, to_kql(format_string))))
 
     def between(self, lower: TimespanType, upper: TimespanType) -> BooleanExpression:
         return BooleanExpression(KQL('{} between ({} .. {})'.format(
@@ -434,6 +472,7 @@ class MappingBaseExpression(BaseExpression):
         return AnyExpression(KQL('{}.{}'.format(self.kql, name)))
 
 
+@plain_expression(List, Tuple)
 class ArrayExpression(ArrayBaseExpression):
     def __len__(self) -> NumberExpression:
         return self.array_length()
@@ -450,12 +489,16 @@ class ArrayExpression(ArrayBaseExpression):
             ', '.join('{}'.format(_subexpr_to_kql(e) for e in elements))
         )))
 
+    def __getitem__(self, index: NumberType) -> BaseExpression:
+        return AnyExpression(KQL('{}[{}]'.format(self.kql, _subexpr_to_kql(index))))
+
     def assign_to(self, *columns: 'Column') -> 'AssignmentBase':
         if len(columns) <= 1:
             return super().assign_to(*columns)
         return AssignmentToMultipleColumns(columns, self)
 
 
+@plain_expression(Mapping)
 class MappingExpression(MappingBaseExpression):
     def keys(self) -> ArrayExpression:
         return ArrayExpression(KQL('bag_keys({})'.format(self.kql)))
@@ -481,6 +524,10 @@ class AnyExpression(
 
 
 class AggregationExpression(BaseExpression):
+    def __new__(cls, *args, **kwargs):
+        if cls is 'AggregationExpression':
+            raise TypeError("AggregationExpression is abstract")
+        return object.__new__(cls)
 
     def assign_to(self, *columns: 'Column') -> 'AssignmentFromAggregationToColumn':
         if len(columns) == 0:
@@ -494,30 +541,37 @@ class AggregationExpression(BaseExpression):
         return self.kql
 
 
+@aggregation_expression(bool)
 class BooleanAggregationExpression(AggregationExpression, BooleanExpression):
     pass
 
 
+@aggregation_expression(Number)
 class NumberAggregationExpression(AggregationExpression, NumberExpression):
     pass
 
 
+@aggregation_expression(str)
 class StringAggregationExpression(AggregationExpression, StringExpression):
     pass
 
 
+@aggregation_expression(datetime)
 class DatetimeAggregationExpression(AggregationExpression, DatetimeExpression):
     pass
 
 
+@aggregation_expression(timedelta)
 class TimespanAggregationExpression(AggregationExpression, TimespanExpression):
     pass
 
 
+@aggregation_expression(List, Tuple)
 class ArrayAggregationExpression(AggregationExpression, ArrayExpression):
     pass
 
 
+@aggregation_expression(Mapping)
 class MappingAggregationExpression(AggregationExpression, MappingExpression):
     pass
 
@@ -561,7 +615,7 @@ class AssignmentToMultipleColumns(AssignmentBase):
 
 
 class AssignmentFromAggregationToColumn(AssignmentBase):
-    def __init__(self, column: Optional['Column'], aggregation: AggregationType) -> None:
+    def __init__(self, column: Optional['Column'], aggregation: AggregationExpression) -> None:
         super().__init__(None if column is None else column.kql, aggregation)
 
 
@@ -597,3 +651,18 @@ class ColumnGenerator:
 # Recommended usage: from pykusto.expressions import column_generator as col
 # TODO: Is there a way to enforce this to be a singleton?
 column_generator = ColumnGenerator()
+
+
+def to_kql(obj: ExpressionType) -> KQL:
+    """
+    Convert the given expression to KQL. If this is a subexpression of a greater expression, neighboring operators might
+    take precedence over operators included in this expression, causing an incorrect evaluation order.
+    If this is a concern, use `_subexpr_to_kql` instead, which will enclose this expression in parentheses if it is
+    a compound expression.
+
+    :param obj: Expression to convert to KQL
+    :return: KQL that represents the given expression
+    """
+    if isinstance(obj, BaseExpression):
+        return obj.kql
+    return kql_converter.for_obj(obj)
diff --git a/pykusto/functions.py b/pykusto/functions.py
index 6894492..6269870 100644
--- a/pykusto/functions.py
+++ b/pykusto/functions.py
@@ -1,7 +1,12 @@
-from pykusto import utils
-from pykusto.expressions import *
-from pykusto.expressions import _subexpr_to_kql, Column
-from pykusto.utils import KQL
+from typing import Union
+
+from pykusto.expressions import Column, NumberType, NumberExpression, TimespanType, \
+    DatetimeExpression, TimespanExpression, ArrayType, DynamicType, DatetimeType, BaseExpression, BooleanType, \
+    ExpressionType, AggregationExpression, StringType, StringExpression, BooleanExpression, \
+    NumberAggregationExpression, MappingAggregationExpression, ArrayAggregationExpression, to_kql, DynamicExpression, \
+    ArrayExpression
+from pykusto.kql_converters import KQL
+from pykusto.type_utils import plain_expression
 
 
 # Scalar functions
@@ -101,9 +106,9 @@ def bin_auto(expr: Union[NumberType, DatetimeType, TimespanType]) -> BaseExpress
 
 
 def case(predicate: BooleanType, val: ExpressionType, *args: Union[BooleanType, ExpressionType]) -> BaseExpression:
-    res = 'case({}, {}, {})'.format(_subexpr_to_kql(predicate),
-                                    _subexpr_to_kql(val),
-                                    ', '.join([_subexpr_to_kql(arg) for arg in args]))
+    res = 'case({}, {}, {})'.format(
+        to_kql(predicate), to_kql(val), ', '.join([to_kql(arg) for arg in args])
+    )
     return AggregationExpression(KQL(res))
 
 
@@ -167,7 +172,7 @@ def cos(expr: NumberType) -> NumberExpression:
 
 
 # def dcount_hll(expr: ExpressionType) -> BaseExpression:
-#     return BaseExpression(KQL('dcount_hll({})'.format(expr)))
+#     return BaseExpression(KQL('dcount_hll({})'.format(to_kql(expr))))
 
 
 # def degrees(self): return
@@ -260,7 +265,12 @@ def hourofday(expr: DatetimeType) -> NumberExpression:
 
 
 def iff(predicate: BooleanType, if_true: ExpressionType, if_false: ExpressionType) -> BaseExpression:
-    return BaseExpression(KQL('iff({}, {}, {})'.format(predicate, _subexpr_to_kql(if_true), _subexpr_to_kql(if_false))))
+    return_type = type(if_true)
+    if type(if_false) is not return_type:
+        raise TypeError("The second and third arguments must be of the same type")
+    return plain_expression.for_type(return_type)(
+        KQL('iff({}, {}, {})'.format(to_kql(predicate), to_kql(if_true), to_kql(if_false)))
+    )
 
 
 def iif(predicate: BooleanType, if_true: ExpressionType, if_false: ExpressionType) -> BaseExpression:
@@ -335,12 +345,12 @@ def make_datetime(year: NumberType,
                   minute: NumberType = None,
                   second: NumberType = None) -> DatetimeExpression:
     res = 'make_datetime({year}, {month}, {day}, {hour}, {minute}, {second})'.format(
-        year=_subexpr_to_kql(year),
-        month=_subexpr_to_kql(month),
-        day=_subexpr_to_kql(day),
-        hour=_subexpr_to_kql(0 if hour is None else hour),
-        minute=_subexpr_to_kql(0 if minute is None else minute),
-        second=_subexpr_to_kql(0 if second is None else second)
+        year=to_kql(year),
+        month=to_kql(month),
+        day=to_kql(day),
+        hour=to_kql(0 if hour is None else hour),
+        minute=to_kql(0 if minute is None else minute),
+        second=to_kql(0 if second is None else second)
     )
     return DatetimeExpression(KQL(res))
 
@@ -365,7 +375,7 @@ def new_guid(): raise NotImplemented  # TODO
 
 def now(offset: TimespanType = None) -> StringExpression:
     if offset:
-        return StringExpression(KQL('now({})'.format(utils.timedelta_to_kql(offset))))
+        return StringExpression(KQL('now({})'.format(to_kql(offset))))
     return StringExpression(KQL('now()'))
 
 
@@ -390,7 +400,7 @@ def pack_dictionary(): raise NotImplemented  # TODO
 
 
 def parse_json(expr: Union[StringType, DynamicType]) -> DynamicExpression:
-    return DynamicExpression(KQL('parse_json({})'.format(expr)))
+    return DynamicExpression(KQL('parse_json({})'.format(to_kql(expr))))
 
 
 # def parse_path(self): return
@@ -418,7 +428,7 @@ def percentrank_tdigest(): raise NotImplemented  # TODO
 
 
 def pow(expr1: NumberType, expr2: NumberType) -> NumberExpression:
-    return NumberExpression(KQL('pow({}, {})'.format(expr1, expr2)))
+    return NumberExpression(KQL('pow({}, {})'.format(to_kql(expr1), to_kql(expr2))))
 
 
 # def radians(self): return
@@ -543,7 +553,7 @@ def round(expr: NumberType, precision: NumberType = None) -> NumberExpression:
 
 
 def sign(expr: NumberType) -> NumberExpression:
-    return NumberExpression(KQL('sign({})'.format(expr)))
+    return NumberExpression(KQL('sign({})'.format(to_kql(expr))))
 
 
 # def sin(self): return
@@ -554,7 +564,7 @@ def split(string: StringType, delimiter: StringType, requested_index: NumberType
 
 
 def sqrt(expr: NumberType) -> NumberExpression:
-    return NumberExpression(KQL('sqrt({})'.format(expr)))
+    return NumberExpression(KQL('sqrt({})'.format(to_kql(expr))))
 
 
 def startofday(expr: DatetimeType, offset: NumberType = None) -> DatetimeExpression:
@@ -574,52 +584,55 @@ def startofyear(expr: DatetimeType, offset: NumberType = None) -> DatetimeExpres
 
 
 def strcat(expr1: StringType, expr2: StringType, *exprs: StringType) -> StringExpression:
-    res = 'strcat({}, {}'.format(_subexpr_to_kql(expr1),
-                                 _subexpr_to_kql(expr2))
+    res = 'strcat({}, {}'.format(to_kql(expr1),
+                                 to_kql(expr2))
     if len(exprs) > 0:
-        res = res + ', ' + ', '.join([_subexpr_to_kql(expr) for expr in exprs])
+        res = res + ', ' + ', '.join([to_kql(expr) for expr in exprs])
     return StringExpression(KQL(res + ')'))
 
 
 def strcat_array(expr: ArrayType, delimiter: StringType) -> StringExpression:
-    return StringExpression(KQL('strcat_array({}, {})'.format(expr, _subexpr_to_kql(delimiter))))
+    return StringExpression(KQL('strcat_array({}, {})'.format(to_kql(expr), to_kql(delimiter))))
 
 
 def strcat_delim(delimiter: StringType, expr1: StringType, expr2: StringType, *exprs: StringType) -> StringExpression:
     res = 'strcat_delim({}, {}, {}'.format(
-        _subexpr_to_kql(delimiter),
-        _subexpr_to_kql(expr1),
-        _subexpr_to_kql(expr2))
+        to_kql(delimiter),
+        to_kql(expr1),
+        to_kql(expr2))
     if len(exprs) > 0:
-        res = res + ', ' + ', '.join([_subexpr_to_kql(expr) for expr in exprs])
+        res = res + ', ' + ', '.join([to_kql(expr) for expr in exprs])
     return StringExpression(KQL(res + ')'))
 
 
 def strcmp(expr1: StringType, expr2: StringType) -> NumberExpression:
-    return NumberExpression(KQL('strcmp({}, {})'.format(expr1, expr2)))
+    return NumberExpression(KQL('strcmp({}, {})'.format(to_kql(expr1), to_kql(expr2))))
 
 
 def string_size(expr: StringType) -> NumberExpression:
-    return NumberExpression(KQL('string_size({})'.format(expr)))
+    return NumberExpression(KQL('string_size({})'.format(to_kql(expr))))
 
 
 def strlen(expr: StringType) -> NumberExpression:
-    return NumberExpression(KQL('strlen({})'.format(expr)))
+    return NumberExpression(KQL('strlen({})'.format(to_kql(expr))))
 
 
 def strrep(expr: StringType,
            multiplier: NumberType,
            delimiter: StringType = None) -> StringExpression:
     if delimiter is None:
-        res = 'strrep({}, {})'.format(expr, multiplier)
+        res = 'strrep({}, {})'.format(to_kql(expr), to_kql(multiplier))
     else:
-        res = 'strrep({}, {}, {})'.format(expr, multiplier, _subexpr_to_kql(delimiter))
+        res = 'strrep({}, {}, {})'.format(to_kql(expr), to_kql(multiplier),
+                                          to_kql(delimiter))
     return StringExpression(KQL((res)))
 
 
 def substring(expr: StringType, start_index: NumberType, length: NumberType = None) -> StringExpression:
     return StringExpression(KQL(
-        ('substring({}, {})' if length is None else 'substring({}, {}, {})').format(expr, start_index, length)
+        ('substring({}, {})' if length is None else 'substring({}, {}, {})').format(
+            to_kql(expr), to_kql(start_index), to_kql(length)
+        )
     ))
 
 
@@ -630,22 +643,22 @@ def substring(expr: StringType, start_index: NumberType, length: NumberType = No
 
 
 def tobool(expr: ExpressionType) -> BooleanExpression:
-    return BooleanExpression(KQL('tobool({})'.format(expr)))
+    return BooleanExpression(KQL('tobool({})'.format(to_kql(expr))))
 
 
 def toboolean(expr: ExpressionType) -> BooleanExpression:
-    return BooleanExpression(KQL('toboolean({})'.format(expr)))
+    return BooleanExpression(KQL('toboolean({})'.format(to_kql(expr))))
 
 
 def todatetime(expr: StringType) -> DatetimeExpression:
-    return DatetimeExpression(KQL('todatetime({})'.format(expr)))
+    return DatetimeExpression(KQL('todatetime({})'.format(to_kql(expr))))
 
 
 def todecimal(): raise NotImplemented  # TODO
 
 
 def todouble(expr: NumberType) -> NumberExpression:
-    return NumberExpression(KQL("todouble({})".format(expr)))
+    return NumberExpression(KQL("todouble({})".format(to_kql(expr))))
 
 
 def todynamic(): raise NotImplemented  # TODO
@@ -732,11 +745,11 @@ def arg_min(*args: ExpressionType) -> AggregationExpression:
 
 
 def avg(expr: ExpressionType) -> NumberAggregationExpression:
-    return NumberAggregationExpression(KQL('avg({})'.format(expr)))
+    return NumberAggregationExpression(KQL('avg({})'.format(to_kql(expr))))
 
 
 def avgif(expr: ExpressionType, predicate: BooleanType) -> NumberAggregationExpression:
-    return NumberAggregationExpression(KQL('avgif({}, {})'.format(expr, predicate)))
+    return NumberAggregationExpression(KQL('avgif({}, {})'.format(to_kql(expr), to_kql(predicate))))
 
 
 # def buildschema(self):
@@ -749,17 +762,19 @@ def count(col: Column = None) -> NumberAggregationExpression:
 
 
 def countif(predicate: BooleanType) -> NumberAggregationExpression:
-    return NumberAggregationExpression(KQL('countif({})'.format(predicate)))
+    return NumberAggregationExpression(KQL('countif({})'.format(to_kql(predicate))))
 
 
 def dcount(expr: ExpressionType, accuracy: NumberType = None) -> NumberAggregationExpression:
     return NumberAggregationExpression(KQL(
-        ('dcount({})' if accuracy is None else 'dcount({}, {})').format(expr, accuracy)
+        ('dcount({})' if accuracy is None else 'dcount({}, {})').format(to_kql(expr), to_kql(accuracy))
     ))
 
 
 def dcountif(expr: ExpressionType, predicate: BooleanType, accuracy: NumberType = 0) -> NumberAggregationExpression:
-    return NumberAggregationExpression(KQL('dcountif({}, {}, {})'.format(expr, predicate, accuracy)))
+    return NumberAggregationExpression(KQL('dcountif({}, {}, {})'.format(
+        to_kql(expr), to_kql(predicate), to_kql(accuracy)
+    )))
 
 
 # def hll(expr: ExpressionType, accuracy: NumberType = None) -> AggregationExpression:
@@ -769,64 +784,64 @@ def dcountif(expr: ExpressionType, predicate: BooleanType, accuracy: NumberType
 
 
 # def hll_merge(expr: ExpressionType) -> AggregationExpression:
-#     return AggregationExpression(KQL('hll_merge({})'.format(expr)))
+#     return AggregationExpression(KQL('hll_merge({})'.format(to_kql(expr))))
 
 
 def make_bag(expr: ExpressionType, max_size: NumberType = None) -> MappingAggregationExpression:
     if max_size:
         return MappingAggregationExpression(KQL('make_bag({}, {})'.format(expr, max_size)))
-    return MappingAggregationExpression(KQL('make_bag({})'.format(expr)))
+    return MappingAggregationExpression(KQL('make_bag({})'.format(to_kql(expr))))
 
 
 def make_list(expr: ExpressionType, max_size: NumberType = None) -> ArrayAggregationExpression:
     if max_size:
         return ArrayAggregationExpression(KQL('make_list({}, {})'.format(expr, max_size)))
-    return ArrayAggregationExpression(KQL('make_list({})'.format(expr)))
+    return ArrayAggregationExpression(KQL('make_list({})'.format(to_kql(expr))))
 
 
 def make_set(expr: ExpressionType, max_size: NumberType = None) -> ArrayAggregationExpression:
     if max_size:
         return ArrayAggregationExpression(KQL('make_set({}, {})'.format(expr, max_size)))
-    return ArrayAggregationExpression(KQL('make_set({})'.format(expr)))
+    return ArrayAggregationExpression(KQL('make_set({})'.format(to_kql(expr))))
 
 
 def max(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('max({})'.format(expr)))
+    return AggregationExpression(KQL('max({})'.format(to_kql(expr))))
 
 
 def min(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('min({})'.format(expr)))
+    return AggregationExpression(KQL('min({})'.format(to_kql(expr))))
 
 
 def percentile(expr: ExpressionType, per: NumberType) -> AggregationExpression:
-    res = 'percentiles({}, {})'.format(expr, _subexpr_to_kql(per))
+    res = 'percentiles({}, {})'.format(expr, to_kql(per))
     return AggregationExpression(KQL(res))
 
 
 def percentiles(expr: ExpressionType, *pers: NumberType) -> AggregationExpression:
     res = 'percentiles({}, {})'.format(expr,
-                                       ', '.join([str(_subexpr_to_kql(per)) for per in pers]))
+                                       ', '.join([str(to_kql(per)) for per in pers]))
     return AggregationExpression(KQL(res))
 
 
 def stdev(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('stdev({})'.format(expr)))
+    return AggregationExpression(KQL('stdev({})'.format(to_kql(expr))))
 
 
 def stdevif(expr: ExpressionType, predicate: BooleanType) -> AggregationExpression:
-    return AggregationExpression(KQL('stdevif({}, {})'.format(expr, predicate)))
+    return AggregationExpression(KQL('stdevif({}, {})'.format(to_kql(expr), to_kql(predicate))))
 
 
 def stdevp(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('stdevp({})'.format(expr)))
+    return AggregationExpression(KQL('stdevp({})'.format(to_kql(expr))))
 
 
 def sum(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('sum({})'.format(expr)))
+    return AggregationExpression(KQL('sum({})'.format(to_kql(expr))))
 
 
 def sumif(expr: ExpressionType, predicate: BooleanType) -> AggregationExpression:
-    return AggregationExpression(KQL('sumif({}, {})'.format(expr, predicate)))
+    return AggregationExpression(KQL('sumif({}, {})'.format(to_kql(expr), to_kql(predicate))))
 
 
 # def tdigest(self):
@@ -838,12 +853,12 @@ def sumif(expr: ExpressionType, predicate: BooleanType) -> AggregationExpression
 
 
 def variance(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('variance({})'.format(expr)))
+    return AggregationExpression(KQL('variance({})'.format(to_kql(expr))))
 
 
 def varianceif(expr: ExpressionType, predicate: BooleanType) -> AggregationExpression:
-    return AggregationExpression(KQL('varianceif({}, {})'.format(expr, predicate)))
+    return AggregationExpression(KQL('varianceif({}, {})'.format(to_kql(expr), to_kql(predicate))))
 
 
 def variancep(expr: ExpressionType) -> AggregationExpression:
-    return AggregationExpression(KQL('variancep({})'.format(expr)))
+    return AggregationExpression(KQL('variancep({})'.format(to_kql(expr))))
diff --git a/pykusto/utils.py b/pykusto/kql_converters.py
similarity index 60%
rename from pykusto/utils.py
rename to pykusto/kql_converters.py
index 521900d..63079f4 100644
--- a/pykusto/utils.py
+++ b/pykusto/kql_converters.py
@@ -1,20 +1,19 @@
 import json
-import logging
 from datetime import datetime, timedelta
-from typing import Union, Mapping, NewType, Type, Dict, Callable, Any, Tuple, List
+from numbers import Number
+from typing import NewType, Mapping, Union, List, Tuple
 
-logger = logging.getLogger("pykusto")
-
-KustoTypes = Union[str, int, bool, datetime, Mapping, List, Tuple, float, timedelta]
-# TODO: Unhandled date types: guid, decimal
+from pykusto.type_utils import kql_converter
 
 KQL = NewType('KQL', str)
 
 
+@kql_converter(datetime)
 def datetime_to_kql(dt: datetime) -> KQL:
     return KQL(dt.strftime('datetime(%Y-%m-%d %H:%M:%S.%f)'))
 
 
+@kql_converter(timedelta)
 def timedelta_to_kql(td: timedelta) -> KQL:
     hours, remainder = divmod(td.seconds, 3600)
     minutes, seconds = divmod(remainder, 60)
@@ -27,6 +26,7 @@ def timedelta_to_kql(td: timedelta) -> KQL:
     ))
 
 
+@kql_converter(Mapping, List, Tuple)
 def dynamic_to_kql(d: Union[Mapping, List, Tuple]) -> KQL:
     query = list(json.dumps(d))
     # Issue #11
@@ -48,29 +48,21 @@ def dynamic_to_kql(d: Union[Mapping, List, Tuple]) -> KQL:
     return KQL("".join(query))
 
 
+@kql_converter(bool)
 def bool_to_kql(b: bool) -> KQL:
     return KQL('true') if b else KQL('false')
 
 
+@kql_converter(str)
 def str_to_kql(s: str) -> KQL:
     return KQL('"{}"'.format(s))
 
 
-KQL_CONVERTER_BY_TYPE: Dict[Type, Callable[[Any], KQL]] = {
-    datetime: datetime_to_kql,
-    timedelta: timedelta_to_kql,
-    Mapping: dynamic_to_kql,
-    List: dynamic_to_kql,
-    Tuple: dynamic_to_kql,
-    bool: bool_to_kql,
-    str: str_to_kql,
-    int: KQL,
-    float: KQL,
-}
-
-
-def to_kql(obj: KustoTypes) -> KQL:
-    for kusto_type, converter in KQL_CONVERTER_BY_TYPE.items():
-        if isinstance(obj, kusto_type):
-            return converter(obj)
-    raise ValueError("No KQL converter found for object {} of type {}".format(obj, type(obj)))
+@kql_converter(Number)
+def number_to_kql(n: Number) -> KQL:
+    return KQL(str(n))
+
+
+@kql_converter(type(None))
+def none_to_kql(n: type(None)) -> KQL:
+    return KQL("")
diff --git a/pykusto/logger.py b/pykusto/logger.py
new file mode 100644
index 0000000..6b2cc66
--- /dev/null
+++ b/pykusto/logger.py
@@ -0,0 +1,3 @@
+import logging
+
+logger = logging.getLogger("pykusto")
diff --git a/pykusto/query.py b/pykusto/query.py
index b7157fb..845517c 100644
--- a/pykusto/query.py
+++ b/pykusto/query.py
@@ -10,9 +10,10 @@ from azure.kusto.data.helpers import dataframe_from_result_table
 from pykusto.client import Table
 from pykusto.expressions import BooleanType, ExpressionType, AggregationExpression, OrderType, \
     StringType, AssignmentBase, AssignmentFromAggregationToColumn, AssignmentToSingleColumn, Column, BaseExpression, \
-    AssignmentFromColumnToColumn
+    AssignmentFromColumnToColumn, AnyExpression, to_kql
+from pykusto.kql_converters import KQL
+from pykusto.logger import logger
 from pykusto.udf import stringify_python_func
-from pykusto.utils import KQL, logger, to_kql
 
 
 class Order(Enum):
@@ -127,7 +128,7 @@ class Query:
         return DistinctQuery(self, columns)
 
     def distinct_all(self):
-        return DistinctQuery(self, (BaseExpression(KQL("*")),))
+        return DistinctQuery(self, (AnyExpression(KQL("*")),))
 
     def extend(self, *args: Union[BaseExpression, AssignmentBase], **kwargs: ExpressionType) -> 'ExtendQuery':
         assignments: List[AssignmentBase] = []
@@ -140,7 +141,7 @@ class Query:
             if isinstance(expression, BaseExpression):
                 assignments.append(expression.assign_to(Column(column_name)))
             else:
-                assignments.append(BaseExpression(to_kql(expression)).assign_to(Column(column_name)))
+                assignments.append(AnyExpression(to_kql(expression)).assign_to(Column(column_name)))
         return ExtendQuery(self, *assignments)
 
     def summarize(self, *args: Union[AggregationExpression, AssignmentFromAggregationToColumn],
diff --git a/pykusto/type_utils.py b/pykusto/type_utils.py
new file mode 100644
index 0000000..ca407cd
--- /dev/null
+++ b/pykusto/type_utils.py
@@ -0,0 +1,65 @@
+from datetime import datetime, timedelta
+from numbers import Number
+from typing import Union, Mapping, Type, Dict, Callable, Tuple, List
+
+KustoTypes = Union[str, Number, bool, datetime, Mapping, List, Tuple, timedelta]
+# TODO: Unhandled data types: guid, decimal
+
+
+class TypeRegistrar:
+    """
+    A factory for annotations that are used to create a mapping between Kusto types and python types and functions.
+    Each annotation must be called with a Kusto type as a parameter. The `for_obj` and `for_type` methods
+    can then be used to retrieve the python type or function corresponding to a given Kusto type.
+    """
+    registry: Dict[Type[KustoTypes], Callable]
+
+    def __init__(self, name: str) -> None:
+        """
+        :param name: Name is used for better logging and clearer errors
+        """
+        self.name = name
+        self.registry = {}
+
+    def __repr__(self) -> str:
+        return self.name
+
+    def __call__(self, *types: Type[KustoTypes]) -> Callable:
+        def inner(wrapped):
+            for t in types:
+                previous = self.registry.setdefault(t, wrapped)
+                if previous is not wrapped:
+                    raise TypeError("{}: type already registered: {}".format(self, t.__name__))
+            return wrapped
+
+        return inner
+
+    def for_obj(self, obj: KustoTypes) -> Callable:
+        """
+        Given an object of Kusto type, retrieve the python type or function associated with the object's type, and call
+        it with the given object as a parameter
+
+        :param obj: An object of Kusto type
+        :return: Associated python object
+        """
+        for registered_type, registered_callable in self.registry.items():
+            if isinstance(obj, registered_type):
+                return registered_callable(obj)
+        raise ValueError("{}: no registered callable for object {} of type {}".format(self, obj, type(obj).__name__))
+
+    def for_type(self, t: Type[KustoTypes]) -> Callable:
+        """
+        Given a Kusto type, retrieve the associated python type or function
+
+        :param t: A Kusto type
+        :return: Associated python object
+        """
+        for registered_type, registered_callable in self.registry.items():
+            if issubclass(t, registered_type):
+                return registered_callable
+        raise ValueError("{}: no registered callable for type {}".format(self, t.__name__))
+
+
+kql_converter = TypeRegistrar("KQL Converter")
+plain_expression = TypeRegistrar("Plain expression")
+aggregation_expression = TypeRegistrar("Aggregation expression")
diff --git a/setup.py b/setup.py
index c2641c0..ca8f191 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='pykusto',
-    version='0.0.8',
+    version='0.0.9',
     packages=find_packages(exclude=['test']),
     url='https://github.com/Azure/pykusto',
     license='MIT License',


Test patch:
diff --git a/test/test_base.py b/test/test_base.py
index 7070cec..dde3185 100644
--- a/test/test_base.py
+++ b/test/test_base.py
@@ -3,7 +3,7 @@ import sys
 from typing import Callable
 from unittest import TestCase
 
-from pykusto.utils import logger
+from pykusto.logger import logger
 
 
 class TestBase(TestCase):
diff --git a/test/test_functions.py b/test/test_functions.py
index 45b380a..3ce0a5c 100644
--- a/test/test_functions.py
+++ b/test/test_functions.py
@@ -40,7 +40,7 @@ class TestFunction(TestBase):
 
     def test_case(self):
         self.assertEqual(
-            " | extend bucket = case((foo <= 3), \"Small\", (foo <= 10), \"Medium\", \"Large\")",
+            " | extend bucket = case(foo <= 3, \"Small\", foo <= 10, \"Medium\", \"Large\")",
             Query().extend(bucket=f.case(col.foo <= 3, "Small", col.foo <= 10, 'Medium', 'Large')).render()
         )
 
@@ -372,13 +372,14 @@ class TestFunction(TestBase):
             Query().extend(f.strcat_delim('-', ',', col.foo)).render()
         )
 
+    @unittest.skip("Enabled after #40 is fixed")
     def test_strcat_array(self):
         self.assertEqual(
             " | where (strcat_array(foo, \",\")) == \"A,B,C\"",
             Query().where(f.strcat_array(col.foo, ',') == 'A,B,C').render()
         )
         self.assertEqual(
-            " | where (strcat_array(['A', 'B', 'C'], \",\")) == \"A,B,C\"",
+            " | where (strcat_array(dynamic([\"A\", \"B\", \"C\"]), \",\")) == \"A,B,C\"",
             Query().where(f.strcat_array(['A', 'B', 'C'], ',') == 'A,B,C').render()
         )
 
@@ -458,7 +459,7 @@ class TestFunction(TestBase):
     #         Query().where(f.todatetime('') > datetime.datetime(2019, 7, 23)).render(),
     #         " | where (startofday(foo)) > datetime(2019-07-23 00:00:00.000000)")
     # def todatetime(expr: StringType) -> DatetimeExpression:
-    #     return DatetimeExpression(KQL('todatetime({})'.format(expr)))
+    #     return DatetimeExpression(KQL('todatetime({})'.format(_subexpr_to_kql(expr))))
 
     def test_todouble(self):
         self.assertEqual(
@@ -488,14 +489,13 @@ class TestFunction(TestBase):
             Query().summarize(f.arg_min(col.foo, col.bar, col.fam)).render()
         )
 
-    @unittest.skip("Re-enable once issue #1 is resolved")
     def test_avg(self):
         self.assertEqual(
             " | summarize avg(foo)",
             Query().summarize(f.avg(col.foo)).render()
         )
         self.assertEqual(
-            " | summarize avg(foo)-5",
+            " | summarize avg(foo) - 5",
             Query().summarize(f.avg(col.foo) - 5).render()
         )
 
diff --git a/test/test_utils.py b/test/test_utils.py
index c3614d8..2b2e7aa 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -1,16 +1,81 @@
-from pykusto import utils
+from pykusto.expressions import to_kql
+from pykusto.type_utils import TypeRegistrar
 from test.test_base import TestBase
 
 
 class TestUtils(TestBase):
     def test_dynamic_to_kql(self):
-        dict ={
+        test_dict = {
             "name": "Alan",
             "age": 21,
             "address": ("NY", 36),
             "pets": ["Libby", "Panda", "]", "["]
         }
         self.assertEqual(
-            "{\"name\": \"Alan\", \"age\": 21, \"address\": (\"NY\", 36), \"pets\": (\"Libby\", \"Panda\", \"]\", \"[\")}",
-            utils.to_kql(dict)
+            "{\"name\": \"Alan\", \"age\": 21, \"address\": (\"NY\", 36), "
+            "\"pets\": (\"Libby\", \"Panda\", \"]\", \"[\")}",
+            to_kql(test_dict)
+        )
+
+    def test_type_registrar_for_type(self):
+        test_annotation = TypeRegistrar("Test annotation")
+
+        @test_annotation(str)
+        def str_annotated(s: str) -> str:
+            return "response to " + s
+
+        self.assertEqual(
+            "response to test for_type",
+            test_annotation.for_type(str)("test for_type")
+        )
+
+    def test_type_registrar_for_obj(self):
+        test_annotation = TypeRegistrar("Test annotation")
+
+        @test_annotation(str)
+        def str_annotated(s: str) -> str:
+            return "response to " + s
+
+        self.assertEqual(
+            "response to test for_obj",
+            test_annotation.for_obj("test for_obj")
+        )
+
+    def test_type_registrar_for_type_not_found(self):
+        test_annotation = TypeRegistrar("Test annotation")
+
+        @test_annotation(str)
+        def str_annotated(s: str) -> str:
+            return "response to " + s
+
+        self.assertRaises(
+            ValueError("Test annotation: no registered callable for type bool"),
+            lambda: test_annotation.for_type(bool)("test for_type")
+        )
+
+    def test_type_registrar_for_obj_not_found(self):
+        test_annotation = TypeRegistrar("Test annotation")
+
+        @test_annotation(str)
+        def str_annotated(s: str) -> str:
+            return "response to " + s
+
+        self.assertRaises(
+            ValueError("Test annotation: no registered callable for object True of type bool"),
+            lambda: test_annotation.for_obj(True)
+        )
+
+    def test_type_registrar_collision(self):
+        test_annotation = TypeRegistrar("Test annotation")
+
+        @test_annotation(str)
+        def str_annotated_1(s: str) -> str:
+            return "response to " + s
+
+        def str_annotated_2(s: str) -> str:
+            return "response to " + s
+
+        self.assertRaises(
+            TypeError("Test annotation: type already registered: str"),
+            lambda: test_annotation(str)(str_annotated_2)
         )