fix(api): map string multiplication with integer to repeat method

ibis-project · Mar 7, 2022 · b205922 · b205922
1 parent 5975b79
commit b205922
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 15 deletions.
diff --git a/ibis/backends/base/sql/alchemy/registry.py b/ibis/backends/base/sql/alchemy/registry.py
@@ -458,7 +458,7 @@ def _sort_key(t, expr):
     ops.TypeOf: unary(sa.func.typeof),
     ops.Literal: _literal,
     ops.ValueList: _value_list,
-    ops.NullLiteral: lambda *args: sa.null(),
+    ops.NullLiteral: lambda *_: sa.null(),
     ops.SimpleCase: _simple_case,
     ops.SearchedCase: _searched_case,
     ops.TableColumn: _table_column,

diff --git a/ibis/backends/base/sql/registry/string.py b/ibis/backends/base/sql/registry/string.py
@@ -92,9 +92,9 @@ def startswith(translator, expr):
 
 
 def endswith(translator, expr):
-    arg, start = expr.op().args
+    arg, end = expr.op().args
 
     arg_formatted = translator.translate(arg)
-    end_formatted = translator.translate(start)
+    end_formatted = translator.translate(end)
 
     return f"{arg_formatted} like concat('%', {end_formatted})"
diff --git a/ibis/backends/sqlite/registry.py b/ibis/backends/sqlite/registry.py
@@ -277,6 +277,6 @@ def _rpad(t, expr):
         ops.StandardDev: toolz.compose(
             sa.func._ibis_sqlite_sqrt, variance_reduction('_ibis_sqlite_var')
         ),
-        ops.RowID: lambda t, expr: sa.literal_column('rowid'),
+        ops.RowID: lambda *_: sa.literal_column('rowid'),
     }
 )
diff --git a/ibis/backends/tests/test_string.py b/ibis/backends/tests/test_string.py
@@ -185,7 +185,17 @@ def test_string_col_is_unicode(backend, alltypes, df):
         param(
             lambda t: t.string_col.repeat(2),
             lambda t: t.string_col * 2,
-            id='repeat',
+            id="repeat_method",
+        ),
+        param(
+            lambda t: 2 * t.string_col,
+            lambda t: 2 * t.string_col,
+            id="repeat_left",
+        ),
+        param(
+            lambda t: t.string_col * 2,
+            lambda t: t.string_col * 2,
+            id="repeat_right",
         ),
         param(
             lambda t: t.string_col.translate('0', 'a'),

diff --git a/ibis/expr/api.py b/ibis/expr/api.py
@@ -3890,8 +3890,8 @@ def _string_getitem(
     'rpad': _rpad,
     '__add__': _string_concat,
     '__radd__': lambda *args: _string_concat(*args[::-1]),
-    '__mul__': mul,
-    '__rmul__': mul,
+    '__mul__': _binop_expr('__mul__', ops.Repeat),
+    '__rmul__': _binop_expr('__rmul__', ops.Repeat),
 }
 
 

diff --git a/ibis/tests/expr/test_value_exprs.py b/ibis/tests/expr/test_value_exprs.py
@@ -694,18 +694,26 @@ def test_chained_comparisons_not_allowed(table):
 
 
 @pytest.mark.parametrize(
-    'operation', [operator.add, operator.mul, operator.truediv, operator.sub]
+    'operation',
+    [operator.add, operator.sub, operator.truediv],
 )
-def test_binop_string_type_error(table, operation):
-    # Strings are not valid for any numeric arithmetic
-    ints = table.d
-    strs = table.g
+@pytest.mark.parametrize(("left", "right"), [("d", "g"), ("g", "d")])
+def test_binop_string_type_error(table, operation, left, right):
+    a = table[left]
+    b = table[right]
 
     with pytest.raises(TypeError):
-        operation(ints, strs)
+        operation(a, b)
 
-    with pytest.raises(TypeError):
-        operation(strs, ints)
+
+@pytest.mark.parametrize(("left", "right"), [("d", "g"), ("g", "d")])
+def test_string_mul(table, left, right):
+    a = table[left]
+    b = table[right]
+
+    expr = a * b
+    assert isinstance(expr, ir.StringColumn)
+    assert isinstance(expr.op(), ops.Repeat)
 
 
 @pytest.mark.parametrize(