Repository: Azure/pykusto

Issue:
iff doesn't support sub expressions
The following syntax is not working:
.extend(timeDelta=iff(col.day-col.day1 == 0, timedelta(DAYS_BACK_DIFF), col.day-col.day1))

Comments:

Patch:
diff --git a/.gitignore b/.gitignore
index 456846f..78112ac 100644
--- a/.gitignore
+++ b/.gitignore
@@ -334,3 +334,4 @@ ASALocalRun/
 /.idea/checkstyle-idea.xml
 *.iml
 /.idea/misc.xml
+/venv
diff --git a/pykusto/functions.py b/pykusto/functions.py
index 8450ef8..210958f 100644
--- a/pykusto/functions.py
+++ b/pykusto/functions.py
@@ -120,6 +120,8 @@ def ceiling(expr: NumberType) -> NumberExpression:
 
 def cos(expr: NumberType) -> NumberExpression:
     return expr.cos()
+
+
 #
 # def cot(self): return
 
@@ -257,12 +259,12 @@ def hourofday(expr: DatetimeType) -> NumberExpression:
     return expr.hourofday()
 
 
-def iif(predicate: BooleanType, if_true: ExpressionType, if_false: ExpressionType) -> BaseExpression:
-    return BaseExpression(KQL('iif({}, {}, {})'.format(predicate, if_true, if_false)))
+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))))
 
 
-def iff(predicate: BooleanType, if_true: ExpressionType, if_false: ExpressionType) -> BaseExpression:
-    return BaseExpression(KQL('iff({}, {}, {})'.format(predicate, if_true, if_false)))
+def iif(predicate: BooleanType, if_true: ExpressionType, if_false: ExpressionType) -> BaseExpression:
+    return iff(predicate, if_true, if_false)
 
 
 #
diff --git a/pykusto/query.py b/pykusto/query.py
index 5025977..b7157fb 100644
--- a/pykusto/query.py
+++ b/pykusto/query.py
@@ -12,7 +12,7 @@ from pykusto.expressions import BooleanType, ExpressionType, AggregationExpressi
     StringType, AssignmentBase, AssignmentFromAggregationToColumn, AssignmentToSingleColumn, Column, BaseExpression, \
     AssignmentFromColumnToColumn
 from pykusto.udf import stringify_python_func
-from pykusto.utils import KQL, logger
+from pykusto.utils import KQL, logger, to_kql
 
 
 class Order(Enum):
@@ -137,7 +137,10 @@ class Query:
             else:
                 assignments.append(arg)
         for column_name, expression in kwargs.items():
-            assignments.append(expression.assign_to(Column(column_name)))
+            if isinstance(expression, BaseExpression):
+                assignments.append(expression.assign_to(Column(column_name)))
+            else:
+                assignments.append(BaseExpression(to_kql(expression)).assign_to(Column(column_name)))
         return ExtendQuery(self, *assignments)
 
     def summarize(self, *args: Union[AggregationExpression, AssignmentFromAggregationToColumn],
diff --git a/setup.py b/setup.py
index 605de64..b31e771 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='pykusto',
-    version='0.0.4',
+    version='0.0.5',
     packages=find_packages(exclude=['test']),
     url='https://github.com/Azure/pykusto',
     license='MIT License',


Test patch:
diff --git a/test/test_expressions.py b/test/test_expressions.py
index 3074ef4..7c42e8a 100644
--- a/test/test_expressions.py
+++ b/test/test_expressions.py
@@ -55,6 +55,16 @@ class TestExpressions(TestBase):
             " | extend ['foo.bar'] = (shoo * 2)",
             Query().extend((col.shoo * 2).assign_to(col.foo.bar)).render(),
         )
+        self.assertEqual(
+            " | extend foo = (shoo * 2)",
+            Query().extend(foo=(col.shoo * 2)).render(),
+        )
+
+    def test_extend_const(self):
+        self.assertEqual(
+            " | extend foo = (5), bar = (\"bar\"), other_col = other",
+            Query().extend(foo=5, bar="bar", other_col=col.other).render(),
+        )
 
     def test_between_timespan(self):
         self.assertEqual(
diff --git a/test/test_functions.py b/test/test_functions.py
index 247142d..55fc998 100644
--- a/test/test_functions.py
+++ b/test/test_functions.py
@@ -7,7 +7,7 @@ from pykusto.query import Query
 from test.test_base import TestBase
 
 
-# TODO dcount_hll, iif
+# TODO dcount_hll
 
 class TestFunction(TestBase):
     def test_acos(self):
@@ -629,3 +629,16 @@ class TestFunction(TestBase):
             " | summarize active_days = dcount(bin(timestamp, time(1.0:0:0.0)))",
             Query().summarize(active_days=f.dcount(f.bin(col.timestamp, datetime.timedelta(1)))).render()
         )
+
+    def test_iff(self):
+        self.assertEqual(
+            " | project foo = (iff(foo > (ago(time(2.0:0:0.0))), time(3.0:0:0.0), time(4.0:0:0.0)))",
+            Query().project(foo=f.iff(col.foo > f.ago(datetime.timedelta(2)), datetime.timedelta(3), datetime.timedelta(4))).render()
+        )
+
+    def test_iif(self):
+        # iif is just an alias to iff
+        self.assertEqual(
+            " | project foo = (iff(foo > (ago(time(2.0:0:0.0))), time(3.0:0:0.0), time(4.0:0:0.0)))",
+            Query().project(foo=f.iif(col.foo > f.ago(datetime.timedelta(2)), datetime.timedelta(3), datetime.timedelta(4))).render()
+        )