Merge pull request #69 from NERC-CEH/feature/FPM-586-remove-base-flag-column

Author: richjam

src/time_stream/base.py

@@ -438,34 +438,29 @@ def get_flag_system(self, name: str) -> type[BitwiseFlag]:
         """
         return self._flag_manager.get_flag_system(name)
 
-    def register_flag_column(self, column_name: str, base: str, flag_system: str) -> None:
+    def register_flag_column(self, column_name: str, flag_system: str) -> None:
         """Mark the specified existing column as a flag column.
 
-        This does not modify the DataFrame; it only records that ``name`` is a flag column associated with the
-        value column ``base``, with values handled by the flag system ``flag_system``.
+        This does not modify the DataFrame; it only records that ``column_name`` is a flag column
+        with values handled by the flag system ``flag_system``.
 
         Args:
             column_name: A column name to mark as a flag column.
-            base: Name of the value/data column this flag column refers to.
             flag_system: The name of the flag system.
         """
-        check_columns_in_dataframe(self.df, [column_name, base])
-        self._flag_manager.register_flag_column(column_name, base, flag_system)
+        check_columns_in_dataframe(self.df, column_name)
+        self._flag_manager.register_flag_column(column_name, flag_system)
 
-    def init_flag_column(
-        self, base: str, flag_system: str, column_name: str | None = None, data: int | Sequence[int] = 0
-    ) -> None:
+    def init_flag_column(self, flag_system: str, column_name: str | None = None, data: int | Sequence[int] = 0) -> None:
         """Add a new column to the TimeFrame DataFrame, setting it as a Flag Column.
 
         Args:
-            base: Name of the value/data column this flag column will refer to.
             flag_system: The name of the flag system.
             column_name: Optional name for the new flag column. If omitted, a name of the
-                    form "{base}__flag__{flag_system}" is used.
+                    form ``__flag__{flag_system}`` is used, with an integer suffix appended if the name
+                    already exists (e.g. ``__flag__CORE_FLAGS__1``).
             data: The default value to populate the flag column with. Can be a scalar or list-like. Defaults to 0.
         """
-        check_columns_in_dataframe(self.df, base)
-
         # Validate that the flag system exists
         self.get_flag_system(flag_system)
 
@@ -477,11 +472,16 @@ def init_flag_column(
 
         # Determine name of flag column
         if not column_name:
-            column_name = f"{base}__flag__{flag_system}"
+            column_name = f"__flag__{flag_system}"
+            if column_name in self.df.columns:
+                col_suffix = 1
+                while f"{column_name}__{col_suffix}" in self.df.columns:
+                    col_suffix += 1
+                column_name = f"{column_name}__{col_suffix}"
 
         # Add and register as a flag column
         self._df = self.df.with_columns(data.alias(column_name))
-        self.register_flag_column(column_name, base, flag_system)
+        self.register_flag_column(column_name, flag_system)
         self._column_metadata.sync()
 
     def get_flag_column(self, flag_column_name: str) -> FlagColumn:
@@ -769,19 +769,18 @@ def infill(
     def select(
         self,
         column_names: str | Sequence[str],
-        include_flag_columns: bool = True,
     ) -> Self:
         """Return a new TimeFrame instance to include only the specified columns.
 
-        By default, this:
+        This:
           - carries over TimeFrame-level metadata,
           - prunes column-level metadata to the kept columns,
           - rebuilds the flag manager to include only kept flag columns.
 
+        Flag columns are not automatically included; name them explicitly if you want them retained.
+
         Args:
             column_names:  Column name(s) to retain in the updated TimeFrame.
-            include_flag_columns: If True, include any registered flag columns whose base is among the
-                                  kept value columns.
 
         Returns:
              New TimeFrame instance with only selected columns.
@@ -797,13 +796,6 @@ def select(
         if self.time_name not in column_names:
             column_names.insert(0, self.time_name)
 
-        # Optionally include associated flag columns
-        if include_flag_columns:
-            for flag_name, flag_column in self._flag_manager.flag_columns.items():
-                # include if its base (value col) is being kept
-                if flag_column.base in column_names:
-                    column_names.append(flag_name)
-
         # Build new frame
         new_df = self.df.select(column_names)
 
@@ -824,9 +816,7 @@ def select(
         # keep only flag columns that survived
         for flag_name, flag_column in self._flag_manager.flag_columns.items():
             if flag_name in column_names:
-                new_flag_manager.register_flag_column(
-                    flag_name, flag_column.base, flag_column.flag_system.system_name()
-                )
+                new_flag_manager.register_flag_column(flag_name, flag_column.flag_system.system_name())
 
         tf._flag_manager = new_flag_manager
         tf._column_metadata.sync()
@@ -868,7 +858,7 @@ def __getitem__(self, key: str | Sequence[str]) -> Self:
         """
         if isinstance(key, str):
             key = [key]
-        return self.select(key, include_flag_columns=False)
+        return self.select(key)
 
     def __str__(self) -> str:
         """Return the string representation of the TimeFrame dataframe."""

src/time_stream/flag_manager.py

@@ -6,7 +6,7 @@
 
 Typical use from within the TimeFrame class:
     1) Register a flag system by name (from a dict or an existing BitwiseFlag subclass).
-    2) Register a DataFrame column as a flag column associated with that system.
+    2) Register a DataFrame column as a flag column linked to that system.
     3) Use FlagColumn.add_flag / remove_flag with Polars expressions to set/clear the flag bits.
 """
 
@@ -31,17 +31,14 @@
 class FlagColumn:
     """Represents a flag column in a TimeFrame.
 
-    A flag column stores bitwise flags governed by a specific flag system. Each flag column is associated with a base
-    data column.
+    A flag column stores bitwise flags governed by a specific flag system.
 
     Attributes:
         name: Name of the flag column in the DataFrame.
-        base: Name of the associated value/data column.
         flag_system: The enum class that defines the available flags and their bit values.
     """
 
     name: str
-    base: str
     flag_system: type[BitwiseFlag]
 
     def add_flag(self, df: pl.DataFrame, flag: int | str, expr: pl.Expr = pl.lit(True)) -> pl.DataFrame:
@@ -90,7 +87,7 @@ def __eq__(self, other: object) -> bool:
         if not isinstance(other, FlagColumn):
             return False
 
-        return self.name == other.name and self.base == other.base and self.flag_system == other.flag_system
+        return self.name == other.name and self.flag_system == other.flag_system
 
     # Make class instances unhashable
     __hash__ = None
@@ -101,7 +98,7 @@ class FlagManager:
 
     This class:
       * registers **flag systems** (bit registries) under a string name;
-      * registers **flag columns** that reference a base data column and a specific flag system.
+      * registers **flag columns** linked to a specific flag system.
     """
 
     def __init__(self) -> None:
@@ -173,22 +170,19 @@ def get_flag_system(self, flag_system_name: str) -> type[BitwiseFlag]:
         except KeyError:
             raise FlagSystemNotFoundError(f"No such flag system: '{flag_system_name}'")
 
-    def register_flag_column(self, name: str, base: str, flag_system_name: str) -> None:
+    def register_flag_column(self, name: str, flag_system_name: str) -> None:
         """Mark the specified existing column as a flag column.
 
         Args:
             name: A column name to mark as a flag column.
-            base: Name of the value/data column this flag column refers to.
             flag_system_name: The name of the flag system.
         """
         flag_column = self._flag_columns.get(name)
         if flag_column:
-            raise FlagSystemError(
-                f"Flag column '{name}' already registered. Base: '{flag_column.base}'; System: '{flag_system_name}'."
-            )
+            raise FlagSystemError(f"Flag column '{name}' already registered. System: '{flag_system_name}'.")
         else:
             flag_system = self.get_flag_system(flag_system_name)
-            flag_column = FlagColumn(name, base, flag_system)
+            flag_column = FlagColumn(name, flag_system)
             self._flag_columns[name] = flag_column
 
     def get_flag_column(self, name: str) -> FlagColumn:
@@ -215,7 +209,7 @@ def copy(self) -> Self:
 
         # register flag columns in the new copy with their associated system names
         for name, flag_column in self._flag_columns.items():
-            out.register_flag_column(name, flag_column.base, flag_column.flag_system.system_name())
+            out.register_flag_column(name, flag_column.flag_system.system_name())
 
         return out
 

tests/time_stream/test_base.py

@@ -132,17 +132,23 @@ def test_select_column_doesnt_mutate_original_tf(self) -> None:
         assert col2_tf == expected
         assert_frame_equal(tf.df, original_df)
 
-    def test_select_column_with_flags(self) -> None:
+    def test_select_column_does_not_auto_include_flags(self) -> None:
+        """Flag columns are not automatically included when selecting a data column."""
         tf = TimeFrame(self.df, time_name="time").with_flag_system("system", {"A": 1, "B": 2, "C": 4})
-        tf.init_flag_column("col1", "system", "flag_col")
-
-        expected = TimeFrame(self.df.select(["time", "col1"]), time_name="time").with_flag_system(
-            "system", {"A": 1, "B": 2, "C": 4}
-        )
-        expected.init_flag_column("col1", "system", "flag_col")
+        tf.init_flag_column("system", "flag_col")
 
         result = tf.select(["col1"])
-        assert result == expected
+        assert "flag_col" not in result.df.columns
+        assert "flag_col" not in result.flag_columns
+
+    def test_select_column_with_explicit_flag_column(self) -> None:
+        """Flag columns are preserved when explicitly included in select."""
+        tf = TimeFrame(self.df, time_name="time").with_flag_system("system", {"A": 1, "B": 2, "C": 4})
+        tf.init_flag_column("system", "flag_col")
+
+        result = tf.select(["col1", "flag_col"])
+        assert "flag_col" in result.df.columns
+        assert "flag_col" in result.flag_columns
 
 
 class TestGetItem:
@@ -370,15 +376,15 @@ def test_init_flag_column_success(self) -> None:
         # Shouldn't have any flag columns to start with
         assert tf._flag_manager._flag_columns == {}
 
-        tf.init_flag_column("value", "quality_control", "flag_column")
-        assert tf._flag_manager._flag_columns == {"flag_column": FlagColumn("flag_column", "value", flag_system)}
+        tf.init_flag_column("quality_control", "flag_column")
+        assert tf._flag_manager._flag_columns == {"flag_column": FlagColumn("flag_column", flag_system)}
 
         assert "flag_column" in tf.flag_columns
 
     def test_init_column_adds_to_metadata(self) -> None:
         """Test that a new flag column gets added to the column metadata dict."""
         tf, flag_system = self.setup_tf()
-        tf.init_flag_column("value", "quality_control", "flag_column")
+        tf.init_flag_column("quality_control", "flag_column")
         expected = {"time": {}, "value": {}, "existing_flags": {}, "flag_column": {}}
         assert tf.column_metadata == expected
 
@@ -389,7 +395,7 @@ def test_init_flag_column_with_single_value(self) -> None:
         # Shouldn't have any flag columns to start with
         assert tf._flag_manager._flag_columns == {}
 
-        tf.init_flag_column("value", "quality_control", "flag_column", 1)
+        tf.init_flag_column("quality_control", "flag_column", 1)
         expected_values = pl.Series("flag_column", [1, 1, 1], dtype=pl.Int64)
         assert_series_equal(tf.df["flag_column"], expected_values)
 
@@ -400,7 +406,7 @@ def test_init_flag_column_with_list_value(self) -> None:
         # Shouldn't have any flag columns to start with
         assert tf._flag_manager._flag_columns == {}
 
-        tf.init_flag_column("value", "quality_control", "flag_column", [1, 2, 4])
+        tf.init_flag_column("quality_control", "flag_column", [1, 2, 4])
         expected_values = pl.Series("flag_column", [1, 2, 4], dtype=pl.Int64)
         assert_series_equal(tf.df["flag_column"], expected_values)
 
@@ -411,21 +417,43 @@ def test_with_no_flag_column_name(self) -> None:
         # Shouldn't have any flag columns to start with
         assert tf._flag_manager._flag_columns == {}
 
-        tf.init_flag_column("value", "quality_control")
-        default_name = "value__flag__quality_control"
+        tf.init_flag_column("quality_control")
+        default_name = "__flag__quality_control"
 
-        assert tf._flag_manager._flag_columns == {default_name: FlagColumn(default_name, "value", flag_system)}
+        assert tf._flag_manager._flag_columns == {default_name: FlagColumn(default_name, flag_system)}
 
         assert default_name in tf.flag_columns
 
     def test_init_no_column_name_adds_to_metadata(self) -> None:
         """Test that a new flag column gets added to the column metadata dict."""
         tf, flag_system = self.setup_tf()
 
-        tf.init_flag_column("value", "quality_control")
-        expected = {"time": {}, "value": {}, "existing_flags": {}, "value__flag__quality_control": {}}
+        tf.init_flag_column("quality_control")
+        expected = {"time": {}, "value": {}, "existing_flags": {}, "__flag__quality_control": {}}
         assert tf.column_metadata == expected
 
+    def test_auto_name_collision_avoidance(self) -> None:
+        """Test that a numeric suffix is appended when the auto-generated name already exists."""
+        tf, flag_system = self.setup_tf()
+
+        tf.init_flag_column("quality_control")
+        tf.init_flag_column("quality_control")
+
+        assert "__flag__quality_control" in tf.flag_columns
+        assert "__flag__quality_control__1" in tf.flag_columns
+
+    def test_auto_name_collision_avoidance_multiple(self) -> None:
+        """Test that the numeric suffix increments correctly across multiple collisions."""
+        tf, flag_system = self.setup_tf()
+
+        tf.init_flag_column("quality_control")
+        tf.init_flag_column("quality_control")
+        tf.init_flag_column("quality_control")
+
+        assert "__flag__quality_control" in tf.flag_columns
+        assert "__flag__quality_control__1" in tf.flag_columns
+        assert "__flag__quality_control__2" in tf.flag_columns
+
 
 class TestTimeSeriesEquality:
     def setUp(self) -> None: