allow user to specify decimal separator as well

wandb · Nov 8, 2024 · 20185af · 20185af
1 parent 1253c1d
commit 20185af
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 16 deletions.
diff --git a/weave_query/tests/test_basic_ops.py b/weave_query/tests/test_basic_ops.py
@@ -94,12 +94,26 @@ def test_string_ops():
     # assert weave.use(foo in foobar) == True # Broken
     # assert weave.use(foobar in foo) == False # Broken
 
-    assert weave.use(make_const_node(weave.types.String(), "123,456,008").parseNumberWithSeparator(",")) == 123456008.0
-    assert weave.use(make_const_node(weave.types.String(), "123,456,008").parseNumberWithSeparator(".")) == None
-    assert weave.use(make_const_node(weave.types.String(), "123,456.008").parseNumberWithSeparator(",")) == 123456.008
-    assert weave.use(make_const_node(weave.types.String(), "123_456_008").parseNumberWithSeparator("_")) == 123456008.0
-    assert weave.use(make_const_node(weave.types.String(), "123 456 008").parseNumberWithSeparator(" ")) == 123456008.0
-    assert weave.use(make_const_node(weave.types.String(), "123.456.008").parseNumberWithSeparator(".")) == 123456008.0
+class TestStringParseNumberWithSeparator:
+    def test_parseNumberWithSeparatorWithOnlyThousandsSeparatorSpecified(self):
+        assert weave.use(make_const_node(weave.types.String(), "123,456,008").parseNumberWithSeparator(thousands_separator=",", decimal_separator=None)) == 123456008.0
+        assert weave.use(make_const_node(weave.types.String(), "123,456,008").parseNumberWithSeparator(thousands_separator=".", decimal_separator=None)) == None
+        assert weave.use(make_const_node(weave.types.String(), "123,456.008").parseNumberWithSeparator(thousands_separator=",", decimal_separator=None)) == 123456.008
+        assert weave.use(make_const_node(weave.types.String(), "123_456_008").parseNumberWithSeparator(thousands_separator="_", decimal_separator=None)) == 123456008.0
+        assert weave.use(make_const_node(weave.types.String(), "123 456 008").parseNumberWithSeparator(thousands_separator=" ", decimal_separator=None)) == 123456008.0
+        assert weave.use(make_const_node(weave.types.String(), "123.456.008").parseNumberWithSeparator(thousands_separator=".", decimal_separator=None)) == 123456008.0
+
+    def test_parseNumberWithSeparatorWithOnlyDecimalSeparatorSpecified(self):
+        assert weave.use(make_const_node(weave.types.String(), "123456.008").parseNumberWithSeparator(thousands_separator=None, decimal_separator=".")) == 123456.008
+        assert weave.use(make_const_node(weave.types.String(), "123456,008").parseNumberWithSeparator(thousands_separator=None, decimal_separator=",")) == 123456.008
+        assert weave.use(make_const_node(weave.types.String(), "123456008").parseNumberWithSeparator(thousands_separator=None, decimal_separator=".")) == 123456008.0
+        assert weave.use(make_const_node(weave.types.String(), "123,456,008").parseNumberWithSeparator(thousands_separator=None, decimal_separator=".")) == None
+
+    def test_parseNumberWithSeparatorWithBothSpecified(self):
+        assert weave.use(make_const_node(weave.types.String(), "123,456.008").parseNumberWithSeparator(thousands_separator=",", decimal_separator=".")) == 123456.008
+        assert weave.use(make_const_node(weave.types.String(), "123.456,008").parseNumberWithSeparator(thousands_separator=".", decimal_separator=",")) == 123456.008
+        assert weave.use(make_const_node(weave.types.String(), "123 456,008").parseNumberWithSeparator(thousands_separator=" ", decimal_separator=",")) == 123456.008
+        assert weave.use(make_const_node(weave.types.String(), "123 456,008").parseNumberWithSeparator(thousands_separator=",", decimal_separator=".")) == None
 
 def test_null_consuming_numbers_ops():
     data = [box.box(1), box.box(None), box.box(2)]

diff --git a/weave_query/weave_query/ops_primitives/string.py b/weave_query/weave_query/ops_primitives/string.py
@@ -203,16 +203,20 @@ def to_number(self):
             return float(self)  # type: ignore
         return None
 
-    @op(
-        name="string-parseNumberWithSeparator",
-        input_type={
-            "self": types.String(), 
-            "separator": types.String()
-        },
-        output_type=types.optional(types.Number())
-    )
-    def parse_number_with_separator(self, separator: str):
-        mNumber = self.replace(separator, '')
+    @op(name="string-parseNumberWithSeparator", output_type=types.optional(types.Number()))
+    def parse_number_with_separator(
+        self, 
+        thousands_separator: typing.Optional[str], 
+        decimal_separator: typing.Optional[str],
+    ):
+        mNumber = self
+
+        if thousands_separator:
+            mNumber = mNumber.replace(thousands_separator, '')
+
+        if decimal_separator:
+            mNumber = mNumber.replace(decimal_separator, '.')
+
         try:
             number = float(mNumber)
         except: