If a DataFrame is split on a column containing NaN values, the rows with NaN values will now be included in the results

pyproject.toml

@@ -39,6 +39,7 @@ Source = "https://git.cpi.imtek.uni-freiburg.de/holgi/sensospot_tools.git"
 [project.optional-dependencies]
 dev = [
     "black",
+    "flit",
     "keyring",
     "pre-commit",
     "ruff",

src/sensospot_tools/__init__.py

@@ -3,7 +3,7 @@
 Some small tools for working with parsed Sensospot data.
 """
 
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 
 from .hdr import normalize, select_hdr_data  # noqa: F401
 from .selection import select, split  # noqa: F401

src/sensospot_tools/selection.py

@@ -31,7 +31,10 @@ def select(
     Returns:
         a copy of the DataFrame that has the value in the column
     """
-    selector = data[column] == value
+    if pandas.isna(value):
+        selector = data[column].isna()
+    else:
+        selector = data[column] == value
     return data.loc[selector].copy()
 
 

tests/test_selection.py

@@ -28,7 +28,7 @@ def test_selection_select(example):
     assert list(result["value"]) == [35, 9]
 
 
-def test_selection_split_one_column(example):
+def test_selection_split_one_column_without_na(example):
     from sensospot_tools.selection import split
 
     result = dict(split(example, "carnivore"))
@@ -38,6 +38,20 @@ def test_selection_split_one_column(example):
     assert list(result[False]["value"]) == [35, 9]
 
 
+def test_selection_split_one_column_with_na(example):
+    import numpy
+    from sensospot_tools.selection import split
+
+    example["carnivore"].iloc[1] = numpy.nan
+
+    result = dict(split(example, "carnivore"))
+
+    assert set(result.keys()) == {False, True, numpy.nan}
+    assert list(result[True]["value"]) == [3, 60]
+    assert list(result[False]["value"]) == [35, 9]
+    assert list(result[numpy.nan]["value"]) == [55]
+
+
 def test_selection_split_multiple_columns(example):
     from sensospot_tools.selection import split