Reimplement snakebids expand using SnakemakeFormatter

- Add allow_missing support to SnakemakeFormatter with constraint preservation
- Store _field_constraints in parse() for use by get_value() with allow_missing
- Add _expand() helper in datasets.py replacing sn_expand calls
- BidsPartialComponent.expand and BidsComponentRow.expand now use _expand()
- None values in extra wildcards converted to empty string
- AnnotatedString flags propagated to expanded outputs (gated on snakemake import)
- Add focused tests for allow_missing, None handling, and AnnotatedString propagation"

Co-authored-by: pvandyken <87136354+pvandyken@users.noreply.github.com>

Author: Copilot

src/snakebids/core/datasets.py

@@ -19,12 +19,95 @@
 from snakebids.exceptions import DuplicateComponentError
 from snakebids.io.console import get_console_size
 from snakebids.io.printing import format_zip_lists, quote_wrap
-from snakebids.snakemake_compat import expand as sn_expand
+from snakebids.snakemake_compat import WildcardError
 from snakebids.types import ZipList
 from snakebids.utils.containers import ImmutableList, MultiSelectDict
 from snakebids.utils.snakemake_templates import SnakemakeFormatter
 from snakebids.utils.utils import get_wildcard_dict, property_alias, zip_list_eq
 
+try:
+    from snakemake.io import AnnotatedString as _AnnotatedString
+
+    def _get_flags(path: Any) -> dict[str, Any] | None:
+        """Extract flags from an AnnotatedString path for preservation during expansion."""
+        if isinstance(path, _AnnotatedString) and path.flags:
+            return dict(path.flags)
+        return None
+
+    def _make_annotated(s: str, flags: dict[str, Any]) -> str:
+        """Create an AnnotatedString with the given flags copied from the template."""
+        result = _AnnotatedString(s)
+        result.flags.update(flags)
+        return result
+
+except ImportError:
+
+    def _get_flags(path: Any) -> dict[str, Any] | None:  # type: ignore[misc]
+        """Return None when snakemake is not installed (no AnnotatedString support)."""
+        return None
+
+    def _make_annotated(s: str, flags: dict[str, Any]) -> str:  # type: ignore[misc]
+        """Return the string unchanged when snakemake is not installed."""
+        return s
+
+
+def _expand(
+    paths: Iterable[Any],
+    zip_lists: ZipList,
+    allow_missing: bool,
+    **wildcards: Any,
+) -> list[str]:
+    """Expand template paths using SnakemakeFormatter.
+
+    For each template path, iterates over rows of the zip-list and optional
+    extra wildcard combinations (product), formatting each path per row.
+    None values in wildcards are converted to "".
+    Output is order-preserving deduplicated.
+    """
+    formatter = SnakemakeFormatter(allow_missing=allow_missing)
+
+    # Normalize extra wildcards: convert None (scalar or in list) → "", ensure lists
+    extra: dict[str, list[str]] = {}
+    for k, vals in wildcards.items():
+        if vals is None:
+            extra[k] = [""]
+        else:
+            extra[k] = ["" if v is None else v for v in itx.always_iterable(vals)]
+
+    # Build zip-list rows
+    if zip_lists:
+        rows: list[dict[str, str]] = [
+            dict(zip(zip_lists.keys(), vals, strict=True))
+            for vals in zip(*zip_lists.values(), strict=True)
+        ]
+    else:
+        rows = [{}]
+
+    # Compute extra wildcard combinations (product)
+    if extra:
+        extra_keys = list(extra.keys())
+        extra_combos: list[tuple[str, ...]] = list(it.product(*extra.values()))
+    else:
+        extra_keys = []
+        extra_combos = [()]
+
+    try:
+        results: list[str] = []
+        for path in paths:
+            path_str = str(path)
+            flags = _get_flags(path)
+            for row in rows:
+                for combo in extra_combos:
+                    kwargs = {**row, **dict(zip(extra_keys, combo, strict=True))}
+                    result: str = formatter.format(path_str, **kwargs)
+                    if flags is not None:
+                        result = _make_annotated(result, flags)
+                    results.append(result)
+    except KeyError as e:
+        raise WildcardError(f"No values given for wildcard {e}.") from e
+
+    return list(dict.fromkeys(results))
+
 
 class BidsDatasetDict(TypedDict):
     """Dict equivalent of BidsInputs, for backwards-compatibility."""
@@ -122,16 +205,11 @@ def expand(
             Keywords not found in the path will be ignored. Keywords take values or
             lists of values to be expanded over the provided paths.
         """
-        return sn_expand(
+        return _expand(
             list(itx.always_iterable(paths)),
-            allow_missing=allow_missing
-            if isinstance(allow_missing, bool)
-            else list(itx.always_iterable(allow_missing)),
-            **{self.entity: list(dict.fromkeys(self._data))},
-            **{
-                wildcard: list(itx.always_iterable(v))
-                for wildcard, v in wildcards.items()
-            },
+            zip_lists={self.entity: list(dict.fromkeys(self._data))},
+            allow_missing=bool(allow_missing),
+            **wildcards,
         )
 
     def filter(
@@ -376,39 +454,18 @@ def expand(
             Keywords not found in the path will be ignored. Keywords take values or
             lists of values to be expanded over the provided paths.
         """
-
-        def sequencify(item: bool | str | Iterable[str]) -> bool | list[str]:
-            if isinstance(item, bool):
-                return item
-            return list(itx.always_iterable(item))
-
-        allow_missing_seq = sequencify(allow_missing)
-        if self.zip_lists:
-            inner_expand = list(
-                # order preserving deduplication
-                dict.fromkeys(
-                    sn_expand(
-                        list(itx.always_iterable(paths)),
-                        zip,
-                        allow_missing=True if wildcards else allow_missing_seq,
-                        **self.zip_lists,
-                    )
-                )
-            )
-        else:
-            inner_expand = list(itx.always_iterable(paths))
-        if not wildcards:
-            return inner_expand
-
-        return sn_expand(
-            inner_expand,
-            allow_missing=allow_missing_seq,
-            # Turn all the wildcard items into lists because Snakemake doesn't handle
-            # iterables very well
-            **{
-                wildcard: list(itx.always_iterable(v))
-                for wildcard, v in wildcards.items()
-            },
+        path_list = list(itx.always_iterable(paths))
+
+        # When zip_lists is empty and no extra wildcards, return paths as-is
+        # (no formatting: avoids errors on arbitrary path text with {wildcards})
+        if not self.zip_lists and not wildcards:
+            return path_list
+
+        return _expand(
+            path_list,
+            zip_lists=self.zip_lists,
+            allow_missing=bool(allow_missing),
+            **wildcards,
         )
 
     def filter(

src/snakebids/utils/snakemake_templates.py

@@ -116,9 +116,11 @@ class SnakemakeFormatter(string.Formatter):
     _UNEXPECTED_OPEN = "unexpected '{' in field name"
 
     @override
-    def __init__(self) -> None:
+    def __init__(self, allow_missing: bool = False) -> None:
         super().__init__()
         self.squelch_underscore = True
+        self.allow_missing = allow_missing
+        self._field_constraints: dict[str, str] = {}
 
     @overload
     def vformat(
@@ -137,6 +139,7 @@ def vformat(
     ) -> str:
         """Call base vformat after resetting squelch_underscore."""
         self.squelch_underscore = True
+        self._field_constraints = {}
         return super().vformat(format_string, args, kwargs)
 
     @override
@@ -214,6 +217,8 @@ def parse(  # noqa: PLR0912
             comma = format_string.find(",", i + 1, j)
             if comma != -1:
                 field_name = format_string[i + 1 : comma]
+                # Store full text (with constraint) for allow_missing support
+                self._field_constraints[field_name] = format_string[i + 1 : close]
             else:
                 field_name = format_string[i + 1 : j]
 
@@ -302,8 +307,12 @@ def get_value(
         elif key.startswith("_") and key.endswith("_") and len(key) > 1:
             entity = key[1:-1]
 
-        # Otherwise, not a special key, so error
+        # Otherwise, not a special key, so error (or preserve if allow_missing)
         else:
+            if self.allow_missing:
+                # Return original brace-wrapped wildcard including constraint
+                original = self._field_constraints.get(key, key)
+                return f"{{{original}}}"
             raise KeyError(key)
 
         # Rule 2.5. Handle directory wildcards

tests/test_datasets.py

@@ -810,3 +810,84 @@ def test_order_of_selectors_is_preserved(
         assert tuple(dicts.zip_lists[selectors]) == tuple(
             itx.unique_everseen(selectors)
         )
+
+
+def _has_annotated_string() -> bool:
+    try:
+        from snakemake.io import AnnotatedString  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+class TestExpandNoneHandling:
+    """Tests for None value conversion in expand()."""
+
+    def test_none_extra_wildcard_treated_as_empty_string(self):
+        """None values in extra wildcards are converted to empty string."""
+        component = BidsComponent(
+            name="test",
+            path="sub-{subject}",
+            zip_lists={"subject": ["001"]},
+        )
+        # Path with optional entity that gets a None value
+        paths = component.expand(
+            "sub-{subject}{_acq_}{acq}",
+            acq=None,
+        )
+        # None → "" → optional entity not included
+        assert paths == ["sub-001"]
+
+    def test_none_in_extra_wildcard_list(self):
+        """None values within a list of extra wildcards are converted to ''."""
+        component = BidsComponent(
+            name="test",
+            path="sub-{subject}",
+            zip_lists={"subject": ["001"]},
+        )
+        paths = component.expand(
+            "sub-{subject}{_acq_}{acq}",
+            acq=["mprage", None],
+        )
+        assert "sub-001_acq-mprage" in paths
+        assert "sub-001" in paths
+
+
+class TestExpandAnnotatedString:
+    """Tests for AnnotatedString flag propagation in expand()."""
+
+    @pytest.mark.skipif(
+        not _has_annotated_string(),
+        reason="snakemake AnnotatedString not available",
+    )
+    def test_annotated_string_flags_propagated_to_outputs(self):
+        """AnnotatedString flags are preserved on expanded output paths."""
+        from snakemake.io import AnnotatedString
+
+        component = BidsComponent(
+            name="test",
+            path="sub-{subject}",
+            zip_lists={"subject": ["001", "002"]},
+        )
+        template = AnnotatedString("sub-{subject}_T1w.nii.gz")
+        template.flags["temp"] = True
+
+        paths = component.expand(template)
+        for path in paths:
+            assert isinstance(path, AnnotatedString)
+            assert path.flags.get("temp") is True
+
+    @pytest.mark.skipif(
+        not _has_annotated_string(),
+        reason="snakemake AnnotatedString not available",
+    )
+    def test_plain_string_not_annotated(self):
+        """Plain string templates produce plain string outputs."""
+        component = BidsComponent(
+            name="test",
+            path="sub-{subject}",
+            zip_lists={"subject": ["001"]},
+        )
+        paths = component.expand("sub-{subject}_T1w.nii.gz")
+        assert type(paths[0]) is str  # noqa: E721

tests/test_snakemake_templates/test_snakemake_formatter.py

@@ -605,3 +605,53 @@ def test_format_complex_path_with_constraints(self):
             acq="mprage",
         )
         assert result == "sub-001_acq-mprage_T1w.nii.gz"
+
+
+class TestAllowMissing:
+    """Tests for allow_missing behavior in SnakemakeFormatter."""
+
+    def test_allow_missing_preserves_plain_wildcard(self):
+        """Missing wildcard with allow_missing=True is preserved as {name}."""
+        formatter = SnakemakeFormatter(allow_missing=True)
+        result = formatter.format("sub-{subject}_{run}_T1w.nii.gz", subject="001")
+        assert result == "sub-001_{run}_T1w.nii.gz"
+
+    def test_allow_missing_preserves_wildcard_with_constraint(self):
+        """Missing wildcard with constraint is preserved with original {name,constraint}."""
+        formatter = SnakemakeFormatter(allow_missing=True)
+        result = formatter.format(
+            r"sub-{subject,\d+}_{run,\d+}_T1w.nii.gz", subject="001"
+        )
+        assert result == r"sub-001_{run,\d+}_T1w.nii.gz"
+
+    def test_allow_missing_false_raises_for_missing_wildcard(self):
+        """Missing wildcard with allow_missing=False raises KeyError."""
+        formatter = SnakemakeFormatter(allow_missing=False)
+        with pytest.raises(KeyError):
+            formatter.format("sub-{subject}_{run}_T1w.nii.gz", subject="001")
+
+    @given(
+        name=safe_field_names(min_size=1).filter(
+            # Exclude numeric names (positional args, not named wildcards)
+            lambda s: not s.isdigit()
+        ),
+        constraint=constraints(),
+    )
+    def test_allow_missing_preserves_constraint_from_template(
+        self, name: str, constraint: str
+    ):
+        """Constraint from template is preserved verbatim in allow_missing output."""
+        formatter = SnakemakeFormatter(allow_missing=True)
+        template = f"{{{name},{constraint}}}"
+        result = formatter.format(template)
+        assert result == template
+
+    def test_allow_missing_resets_between_vformat_calls(self):
+        """_field_constraints is reset on each vformat call."""
+        formatter = SnakemakeFormatter(allow_missing=True)
+        # First call with constraint
+        r1 = formatter.format(r"sub-{subject,\d+}")
+        assert r1 == r"sub-{subject,\d+}"
+        # Second call without constraint: should use the new (no-constraint) form
+        r2 = formatter.format("{subject}")
+        assert r2 == "{subject}"