Correctly handle null column descriptions

Author: AstroPatty

src/opencosmo/dataset/dataset.py

@@ -737,6 +737,7 @@ def with_new_columns(
         """
         if isinstance(descriptions, str):
             descriptions = {key: descriptions for key in new_columns.keys()}
+        print(descriptions)
         new_state = self.__state.with_new_columns(descriptions, **new_columns)
         return Dataset(self.__handler, self.__header, new_state, self.__tree)
 

src/opencosmo/dataset/state.py

@@ -197,6 +197,7 @@ def with_new_columns(
         new_im_handler = self.__im_handler
         derived_update = {}
         new_unit_handler = self.__unit_handler
+        print(new_columns)
         for name, new_col in new_columns.items():
             if name in column_names:
                 raise ValueError(f"Dataset already has column named {name}")
@@ -216,7 +217,7 @@ def with_new_columns(
                     new_unit_handler = new_unit_handler.with_new_columns(**{name: None})
 
                 new_im_handler = new_im_handler.with_new_column(
-                    name, new_col, descriptions.get(name)
+                    name, new_col, descriptions.get(name, "None")
                 )
 
             elif not new_col.check_parent_existance(column_names):
@@ -226,7 +227,7 @@ def with_new_columns(
                 )
             else:
                 unit = new_col.get_units(self.__unit_handler.base_units)
-                new_col.description = descriptions.get(name)
+                new_col.description = descriptions.get(name, "None")
                 new_unit_handler = new_unit_handler.with_new_columns(**{name: unit})
                 derived_update[name] = new_col
             column_names.add(name)

test/test_dataset.py

@@ -60,19 +60,19 @@ def test_descriptions_after_insert(input_path, tmp_path):
     p = tmp_path / "test.hdf5"
     random_data = np.random.randint(0, 100, len(ds))
     ds = ds.with_new_columns(random_data=random_data)
-    assert ds.descriptions["random_data"] is None
+    assert ds.descriptions["random_data"] == "None"
     oc.write(p, ds)
     ds = oc.open(p)
     descriptions = ds.descriptions
-    assert descriptions["random_data"] is None
+    assert descriptions["random_data"] == "None"
 
 
 def test_description_with_insert(input_path, tmp_path):
     ds = oc.open(input_path)
     p = tmp_path / "test.hdf5"
     random_data = np.random.randint(0, 100, len(ds))
     ds = ds.with_new_columns(
-        random_data=random_data, description="random data for a test"
+        random_data=random_data, descriptions="random data for a test"
     )
     assert ds.descriptions["random_data"] == "random data for a test"
     oc.write(p, ds)
@@ -89,7 +89,7 @@ def test_description_with_insert_multiple(input_path, tmp_path):
     ds = ds.with_new_columns(
         random_data=random_data,
         halo_px=fof_halo_px,
-        description={
+        descriptions={
             "random_data": "random data for a test",
             "halo_px": "halo x momentum",
         },