Merge pull request #32 from jasondet/main

Avoid `nda.tolist()` in `Table.get_dataframe()` when possible
legend-exp · Nov 3, 2023 · 1607492 · 1607492
2 parents 763769d + e53fc41
commit 1607492
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/lgdo/types/scalar.py b/src/lgdo/types/scalar.py
@@ -18,7 +18,7 @@ class Scalar(LGDO):
 
     # TODO: do scalars need proper numpy dtypes?
 
-    def __init__(self, value: int | float, attrs: dict[str, Any] = None) -> None:
+    def __init__(self, value: int | float | str, attrs: dict[str, Any] = None) -> None:
         """
         Parameters
         ----------

diff --git a/src/lgdo/types/table.py b/src/lgdo/types/table.py
@@ -225,7 +225,10 @@ def get_dataframe(
                 if not hasattr(column, "nda"):
                     raise ValueError(f"column {col} does not have an nda")
                 else:
-                    df[prefix + str(col)] = column.nda.tolist()
+                    if len(column.nda.shape) == 1:
+                        df[prefix + str(col)] = column.nda
+                    else:
+                        df[prefix + str(col)] = column.nda.tolist()
 
         return df