Improve hamilton-ui's dataframe_stats compatibility with numpy v2

DAGWorks-Inc · Aug 9, 2024 · 42f5ed7 · 42f5ed7
1 parent 75a11dc
commit 42f5ed7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/ui/sdk/src/hamilton_sdk/tracking/dataframe_stats.py b/ui/sdk/src/hamilton_sdk/tracking/dataframe_stats.py
@@ -13,7 +13,7 @@ def type_converter(obj: Any) -> Any:
         result = int(obj)
     elif isinstance(obj, np.floating):
         result = float(obj)
-    elif isinstance(obj, np.complex_):
+    elif isinstance(obj, np.complexfloating) or isinstance(obj, getattr(np, "complex_", ())):
         result = complex(obj)
     elif isinstance(obj, dict):
         result = {}

diff --git a/ui/sdk/tests/tracking/test_dataframe_stats.py b/ui/sdk/tests/tracking/test_dataframe_stats.py
@@ -3,7 +3,13 @@
 import numpy as np
 import pandas as pd
 from hamilton_sdk.tracking import dataframe_stats
-from pytest import mark
+from pytest import mark, param
+
+
+skip_NAN_on_numpy_v2 = mark.skipif(
+    not hasattr(np, "NAN"),
+    reason="NAN is not available in numpy v2",
+)
 
 
 # Tests the type converter
@@ -21,7 +27,7 @@
         ({"a": 1, "b": 2, "c": {"d": np.int8(3), "e": 4}}, {"a": 1, "b": 2, "c": {"d": 3, "e": 4}}),
         (pd.NaT, None),
         (pd.NA, None),
-        (np.NAN, None),
+        param(getattr(np, "NAN", np.nan), None, marks=skip_NAN_on_numpy_v2),
         (np.nan, None),
         (math.nan, None),
         # (math.inf, None),