Added support for PyKX input types

hablapps · Mar 11, 2024 · 4d07e45 · 4d07e45
1 parent 60747e4
commit 4d07e45
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 5 deletions.
diff --git a/src/pykx/pandas_api/pandas_meta.py b/src/pykx/pandas_api/pandas_meta.py
@@ -1,5 +1,3 @@
-from typing import Dict, Union
-
 from . import api_return
 from ..exceptions import QError
 
@@ -268,7 +266,7 @@ def abs(self, numeric_only=False):
         return q.abs(tab)
 
     @api_return
-    def round(self, decimals: Union[int, Dict[str, int]] = 0):
+    def round(self, decimals=0):
         tab = self
         if 'Keyed' in str(type(tab)):
             tab = q.value(tab)
@@ -278,10 +276,13 @@ def round(self, decimals: Union[int, Dict[str, int]] = 0):
 
         cast_back = q('{string[y][0]$x}')
 
-        if isinstance(decimals, int):
+        if isinstance(decimals, int) or q("{-7h~type x}", decimals):
             dec_dict = {col: decimals for col in affected_cols}
-        else:
+        elif isinstance(decimals, dict) or (q("{99h~type x}", decimals) and
+                                            'Keyed' not in str(type(decimals))):
             dec_dict = {col: decimals[col] for col in affected_cols}
+        else:
+            raise TypeError('Parameter "decimals" should be integer or dictionary.')
 
         rounded = {col: [cast_back(round(elem, dec_dict[col]), type_dict[col])
                          for elem in tab[col]]

diff --git a/tests/test_pandas_api.py b/tests/test_pandas_api.py
@@ -1874,6 +1874,31 @@ def test_pandas_round(kx, q):
 
     pd.testing.assert_frame_equal(q_tab.dtypes.pd(), q_res.dtypes.pd())
 
+    q_res = q_tab.round(kx.toq(dict_test))
+    pd.testing.assert_frame_equal(p_tab.round(dict_test), q_res.pd())
+
+    pd.testing.assert_frame_equal(q_tab.dtypes.pd(), q_res.dtypes.pd())
+
+    with pytest.raises(TypeError):
+        q_tab.round(.1)
+
+    err_tab = pd.DataFrame({
+        "time": [
+            pd.Timestamp("2016-05-25 13:30:00.023"),
+            pd.Timestamp("2016-05-25 13:30:00.038"),
+            pd.Timestamp("2016-05-25 13:30:00.048"),
+            pd.Timestamp("2016-05-25 13:30:00.048"),
+            pd.Timestamp("2016-05-25 13:30:00.048")
+        ],
+        "ticker": ["MSFT", "MSFT", "GOOG", "GOOG", "AAPL"],
+        "price": [51.95, 51.95, 720.77, 720.92, 98.0],
+        "quantity": [75, 155, 100, 100, 100]
+    })
+
+    q_err_tab = q('1!', kx.toq(err_tab))
+    with pytest.raises(TypeError):
+        q_tab.round(q_err_tab)
+
 
 def test_pandas_min(q):
     tab = q('([] sym: 100?`foo`bar`baz`qux; price: 250.0f - 100?500.0f; ints: 100 - 100?200)')