Enhance task column cardinality with type groups

DashAI/back/tasks/base_task.py

@@ -41,13 +41,163 @@ def schema(self) -> Dict[str, Any]:
         """
         raise NotImplementedError
 
+    @classmethod
+    def _column_groups(cls, side: str) -> List[Dict[str, Any]]:
+        """Read one side of the contract as a list of type groups.
+
+        A task says what its columns may be as a list of groups, each naming a
+        set of interchangeable types and how many columns of that set it takes:
+
+        .. code-block:: python
+
+            "inputs": [
+                {"types": [Date], "cardinality": 1},
+                {"types": [Float, Integer], "cardinality": {"min": 1, "max": "n"}},
+            ]
+
+        Grouping is what lets a task ask for one date column and any number of
+        numeric ones at the same time. A single cardinality for the whole side
+        cannot say that: it can only demand a total, so it would accept two
+        dates and no numbers as readily as the intended shape.
+
+        The older flat spelling, ``inputs_types`` alongside
+        ``inputs_cardinality``, is still read and means a single group, so
+        tasks and plugins written against it keep working unchanged.
+
+        Parameters
+        ----------
+        side : str
+            Either ``"inputs"`` or ``"outputs"``.
+
+        Returns
+        -------
+        list of dict
+            One entry per group, each with ``"types"`` (a tuple of type
+            classes), ``"min"`` (an int) and ``"max"`` (an int or ``"n"``).
+        """
+        metadata = cls.metadata
+
+        if side in metadata:
+            declared = metadata[side]
+        else:
+            declared = [
+                {
+                    "types": metadata[f"{side}_types"],
+                    "cardinality": metadata[f"{side}_cardinality"],
+                }
+            ]
+
+        return [
+            {
+                "types": tuple(group["types"]),
+                **cls._bounds(group.get("cardinality", "n")),
+            }
+            for group in declared
+        ]
+
+    @staticmethod
+    def _bounds(cardinality: Any) -> Dict[str, Any]:
+        """Read a declared cardinality as a minimum and a maximum.
+
+        Parameters
+        ----------
+        cardinality : Any
+            An int for an exact count, ``"n"`` for any number, or a mapping
+            with ``"min"`` and ``"max"`` for a range whose ``"max"`` may itself
+            be ``"n"``.
+
+        Returns
+        -------
+        dict
+            A mapping with keys ``"min"`` and ``"max"``.
+        """
+        if isinstance(cardinality, dict):
+            return {
+                "min": cardinality.get("min", 0),
+                "max": cardinality.get("max", "n"),
+            }
+        if cardinality == "n":
+            return {"min": 0, "max": "n"}
+        return {"min": cardinality, "max": cardinality}
+
+    @staticmethod
+    def _total_cardinality(groups: List[Dict[str, Any]]) -> Any:
+        """State a whole side's cardinality the way the flat contract did.
+
+        Kept so every consumer of the two-key metadata, the column picker
+        among them, reads the same value it always did for the tasks that
+        declare a single group.
+
+        Parameters
+        ----------
+        groups : list of dict
+            The normalised groups of one side.
+
+        Returns
+        -------
+        int or str
+            The exact total when every group is exact, otherwise ``"n"``.
+        """
+        if all(group["min"] == group["max"] for group in groups):
+            return sum(group["min"] for group in groups)
+        return "n"
+
+    @staticmethod
+    def _type_name(dashai_type) -> str:
+        """Name a type the way the frontend and the columns themselves do.
+
+        A DashAI type reports its own name through ``display_name()``, which
+        matches what a column emits via ``to_string()``. These lists may also
+        hold foreign classes, for example HuggingFace dataset features or a
+        plugin's own types, so those fall back to the class name.
+
+        Parameters
+        ----------
+        dashai_type : type
+            One of the declared type classes.
+
+        Returns
+        -------
+        str
+            The name to show for it.
+        """
+        getter = getattr(dashai_type, "display_name", None)
+        return getter() if callable(getter) else dashai_type.__name__
+
+    @staticmethod
+    def _cardinality_text(group: Dict[str, Any]) -> str:
+        """Say a group's cardinality the way an error message should read it.
+
+        Parameters
+        ----------
+        group : dict
+            A normalised group.
+
+        Returns
+        -------
+        str
+            For example ``"1"``, ``"n"``, ``"at least 1"`` or ``"1 to 3"``.
+        """
+        minimum, maximum = group["min"], group["max"]
+        if maximum == "n":
+            return "n" if minimum == 0 else f"at least {minimum}"
+        if minimum == maximum:
+            return str(minimum)
+        return f"{minimum} to {maximum}"
+
     @classmethod
     def get_metadata(cls) -> Dict[str, Any]:
         """Return serialisable metadata for the current task.
 
-        Converts the ``inputs_types`` and ``outputs_types`` entries from class
-        objects to their string names so the result can be JSON-serialised by
-        the DashAI frontend.
+        The contract is reported twice. ``"inputs"`` and ``"outputs"`` carry
+        the per-group form, which is the only one that can tell "one date and
+        any number of numbers" apart from "any number of dates or numbers".
+        ``"inputs_types"`` and the three keys beside it carry the flattened
+        view every existing consumer already reads, unchanged for the tasks
+        that declare a single group.
+
+        Type classes become their string names throughout so the result can be
+        JSON-serialised by the DashAI frontend.
 
         Parameters
         ----------
@@ -58,32 +208,99 @@ def get_metadata(cls) -> Dict[str, Any]:
         -------
         Dict[str, Any]
             Dictionary with keys ``"inputs_types"``, ``"outputs_types"``,
-            ``"inputs_cardinality"``, and ``"outputs_cardinality"``.
+            ``"inputs_cardinality"``, ``"outputs_cardinality"``, ``"inputs"``
+            and ``"outputs"``.
         """
-        metadata = cls.metadata
+        parsed_metadata: dict = {}
+
+        for side in ("inputs", "outputs"):
+            groups = cls._column_groups(side)
+
+            names: List[str] = []
+            for group in groups:
+                for dashai_type in group["types"]:
+                    name = cls._type_name(dashai_type)
+                    if name not in names:
+                        names.append(name)
+
+            parsed_metadata[f"{side}_types"] = names
+            parsed_metadata[f"{side}_cardinality"] = cls._total_cardinality(groups)
+            parsed_metadata[side] = [
+                {
+                    "types": [cls._type_name(t) for t in group["types"]],
+                    "min": group["min"],
+                    "max": group["max"],
+                }
+                for group in groups
+            ]
 
-        def _name(dashai_type) -> str:
-            """Name the frontend uses, tolerating types from outside DashAI.
-
-            A DashAI type reports its own name through ``display_name()``, which
-            matches what a column emits via ``to_string()``. These lists may
-            also hold foreign classes, for example HuggingFace dataset features
-            or a plugin's own types, so those fall back to the class name.
-            """
-            getter = getattr(dashai_type, "display_name", None)
-            return getter() if callable(getter) else dashai_type.__name__
-
-        inputs_types = [_name(t) for t in metadata["inputs_types"]]
-        outputs_types = [_name(t) for t in metadata["outputs_types"]]
-
-        parsed_metadata: dict = {
-            "inputs_types": inputs_types,
-            "outputs_types": outputs_types,
-            "inputs_cardinality": metadata["inputs_cardinality"],
-            "outputs_cardinality": metadata["outputs_cardinality"],
-        }
         return parsed_metadata
 
+    def _validate_side(
+        self,
+        columns: List[str],
+        types: Dict[str, Any],
+        side: str,
+    ) -> None:
+        """Check one side's columns against the groups the task declares.
+
+        Each column is charged to the first group that accepts its type and
+        still has room, so a task asking for one date and any number of
+        numbers reads a date, a number and a number as a full match rather
+        than as three columns competing for one slot.
+
+        Parameters
+        ----------
+        columns : list of str
+            The selected column names.
+        types : dict
+            The dataset's column types, keyed by column name.
+        side : str
+            Either ``"inputs"`` or ``"outputs"``.
+
+        Raises
+        ------
+        TypeError
+            If a column's type belongs to none of the groups.
+        ValueError
+            If a group ends up with a number of columns outside its bounds.
+        """
+        groups = self._column_groups(side)
+        counts = [0] * len(groups)
+        label = "Input" if side == "inputs" else "Output"
+
+        for column in columns:
+            column_type = types[column]
+            matching = [
+                index
+                for index, group in enumerate(groups)
+                if isinstance(column_type, group["types"])
+            ]
+            if not matching:
+                raise TypeError(
+                    f"{column_type} is not an allowed type for {side[:-1]} columns."
+                )
+
+            with_room = [
+                index
+                for index in matching
+                if groups[index]["max"] == "n" or counts[index] < groups[index]["max"]
+            ]
+            counts[(with_room or matching)[0]] += 1
+
+        for group, count in zip(groups, counts, strict=True):
+            if count < group["min"] or (group["max"] != "n" and count > group["max"]):
+                of_types = (
+                    ""
+                    if len(groups) == 1
+                    else " for columns of type "
+                    + ", ".join(self._type_name(t) for t in group["types"])
+                )
+                raise ValueError(
+                    f"{label} cardinality ({count}) does not match task "
+                    f"cardinality ({self._cardinality_text(group)}){of_types}"
+                )
+
     def validate_dataset_for_task(
         self,
         dataset: "DashAIDataset",
@@ -99,45 +316,14 @@ def validate_dataset_for_task(
             Dataset to be validated
         dataset_name : str
             Dataset name
+        input_columns : list of str
+            Names of the columns selected as inputs.
+        output_columns : list of str
+            Names of the columns selected as outputs.
         """
-        metadata = self.metadata
-        allowed_input_types = tuple(metadata["inputs_types"])
-        allowed_output_types = tuple(metadata["outputs_types"])
-        inputs_cardinality = metadata["inputs_cardinality"]
-        outputs_cardinality = metadata["outputs_cardinality"]
         types = dataset._types
-        # Check input types
-        for input_col in input_columns:
-            input_col_type = types[input_col]
-
-            if not isinstance(input_col_type, allowed_input_types):
-                raise TypeError(
-                    f"{input_col_type} is not an allowed type for input columns."
-                )
-
-        # Check output types
-        for output_col in output_columns:
-            output_col_type = types[output_col]
-
-            if not isinstance(output_col_type, allowed_output_types):
-                raise TypeError(
-                    f"{output_col_type} is not an allowed type for output columns."
-                )
-
-        # Check input cardinality
-        if inputs_cardinality != "n" and len(input_columns) != inputs_cardinality:
-            raise ValueError(
-                f"Input cardinality ({len(input_columns)}) does not"
-                f" match task cardinality ({inputs_cardinality})"
-            )
-
-        # Check output cardinality
-        if outputs_cardinality != "n" and len(output_columns) != outputs_cardinality:
-            raise ValueError(
-                f"Output cardinality ({len(output_columns)})"
-                f" does not "
-                f"match task cardinality ({outputs_cardinality})"
-            )
+        self._validate_side(input_columns, types, "inputs")
+        self._validate_side(output_columns, types, "outputs")
 
     def prepare_for_task(
         self,