Merge USER_CONTENT from add_to_package with user-code walker

USER_CONTENT tuples produced by FlowDecorator/StepDecorator/
FlowMutator/StepMutator add_to_package hooks were previously dropped
on the floor. They are now recorded per-arcname and merged into the
packaged user code: files already yielded by the flow-directory walker
are skipped (dedup), and anything outside the flow dir or filtered out
by the suffix/user filter still makes it into the package.

Also add an example `@package_suffixes` FlowMutator that declares
additional file suffixes to include in the package — mirroring the
`--package-suffixes` CLI option as a decorator and demonstrating the
new USER_CONTENT merging behavior.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: romain-intel

metaflow/package/__init__.py

@@ -132,6 +132,12 @@ def _module_selector(m) -> bool:
         self._blob_url = None
         self._blob = None
 
+        # USER_CONTENT files contributed by decorators/mutators via add_to_package.
+        # Keyed by arcname (path relative to the flow directory) -> absolute file path.
+        # These are merged with the files discovered by walking the flow directory;
+        # duplicates (same arcname) are dropped to avoid conflicts.
+        self._user_content_from_addl: Dict[str, str] = {}
+
         # Update package in the system context -- it will be available
         # in all hooks going forward including the ones called in the
         # thread that is used to create the package asynchronously.
@@ -612,8 +618,22 @@ def _check_tuple(path_tuple):
                         file_name, [deco_module_paths[file_name], file_path]
                     )
             elif file_type == ContentType.USER_CONTENT:
-                pass  # For now pass -- it will be included normally but we need to
-                # merge it.
+                # USER_CONTENT files will be merged with the files discovered by
+                # walking the flow directory. Track them here so we can:
+                #   1. Include them in the package even if they live outside the
+                #      flow directory (or are excluded by the user_code_filter).
+                #   2. Avoid duplicating files already picked up by the walker.
+                real_path = os.path.realpath(path_tuple[0])
+                path_tuple = (real_path, file_name, file_type)
+                existing = self._user_content_from_addl.get(file_name)
+                if existing is None:
+                    self._user_content_from_addl[file_name] = real_path
+                elif existing != real_path:
+                    raise NonUniqueFileNameToFilePathMappingException(
+                        file_name, [existing, real_path]
+                    )
+                else:
+                    return None  # Already recorded for this arcname
             else:
                 raise ValueError(f"Unknown file type: {file_type}")
             return path_tuple
@@ -675,15 +695,20 @@ def _add_tuple(path_tuple):
                     _add_tuple(path_tuple)
 
     def _user_code_tuples(self):
+        # Track arcnames yielded by the directory walker so we can detect overlap
+        # with USER_CONTENT files contributed via add_to_package hooks.
+        seen_arcnames = set()
         if R.use_r():
             # the R working directory
             self._user_flow_dir = R.working_dir()
             for path_tuple in walk(
                 "%s/" % R.working_dir(), file_filter=self._user_code_filter
             ):
+                seen_arcnames.add(path_tuple[1])
                 yield path_tuple
             # the R package
             for path_tuple in R.package_paths():
+                seen_arcnames.add(path_tuple[1])
                 yield path_tuple
         else:
             # the user's working directory
@@ -694,10 +719,17 @@ def _user_code_tuples(self):
                 file_filter=self._user_code_filter,
                 exclude_tl_dirs=self._exclude_tl_dirs,
             ):
-                # TODO: This is where we will check if the file is already included
-                # in the mfcontent portion
+                seen_arcnames.add(path_tuple[1])
                 yield path_tuple
 
+        # Emit USER_CONTENT files contributed by decorators/mutators that were not
+        # already picked up by the directory walker (either because they live
+        # outside the flow directory or were filtered out by the suffix/user filter).
+        for arcname, file_path in self._user_content_from_addl.items():
+            if arcname in seen_arcnames:
+                continue
+            yield (file_path, arcname)
+
     def _make(self):
         backend = self._backend()
         with backend.create() as archive:

metaflow/plugins/package_suffixes_mutator.py

@@ -0,0 +1,69 @@
+"""
+Example FlowMutator that extends the set of file suffixes included in the code
+package — mirrors the ``--package-suffixes`` CLI option as a decorator.
+
+By default Metaflow's packaging walks the flow directory and includes files
+whose suffixes are in ``DEFAULT_PACKAGE_SUFFIXES`` (``.py,.R,.RDS``) plus
+whatever was passed via ``--package-suffixes``. This mutator lets a flow
+author declare additional suffixes directly on the flow class::
+
+    from metaflow import FlowSpec, step
+    from metaflow.plugins.package_suffixes_mutator import package_suffixes
+
+    @package_suffixes([".yaml", ".json"])
+    class MyFlow(FlowSpec):
+        @step
+        def start(self):
+            ...
+
+The mutator walks the flow directory itself and yields every file with a
+matching suffix as ``USER_CONTENT``. The packaging layer deduplicates against
+the files that the default walker already picked up, so files that are already
+included (e.g. ``.py`` files) are not added twice.
+"""
+
+import inspect
+import os
+from typing import List, Optional, Union
+
+from metaflow.packaging_sys import ContentType
+from metaflow.packaging_sys.utils import walk
+from metaflow.user_decorators.user_flow_decorator import FlowMutator
+
+
+class package_suffixes(FlowMutator):
+    """Include additional file suffixes in the code package.
+
+    Parameters
+    ----------
+    suffixes : list of str or comma-separated str
+        Additional file suffixes to include (e.g. ``[".yaml", ".json"]`` or
+        ``".yaml,.json"``). Leading dots are optional; a suffix without a
+        leading dot is treated as an extension (``yaml`` → ``.yaml``).
+    """
+
+    def init(self, suffixes: Union[List[str], str]):
+        if isinstance(suffixes, str):
+            suffixes = [s.strip() for s in suffixes.split(",") if s.strip()]
+        self._suffixes = tuple(
+            (s if s.startswith(".") else "." + s).lower() for s in suffixes
+        )
+
+    def add_to_package(self):
+        if not self._suffixes:
+            return
+
+        try:
+            flow_file = inspect.getfile(self._flow_cls)
+        except (TypeError, OSError):
+            return
+        flow_dir = os.path.dirname(os.path.abspath(flow_file))
+        if not flow_dir or not os.path.isdir(flow_dir):
+            return
+
+        def _filter(fname: str) -> bool:
+            lname = fname.lower()
+            return any(lname.endswith(sfx) for sfx in self._suffixes)
+
+        for file_path, arcname in walk(flow_dir + os.sep, file_filter=_filter):
+            yield (file_path, arcname, ContentType.USER_CONTENT)

test/unit/test_add_to_package.py

@@ -57,12 +57,19 @@ def _make_deco(tuples):
     return deco
 
 
-def _call_add_addl_files(flow, environment, mfcontent):
-    """Call _add_addl_files on a bare MetaflowPackage instance."""
+def _build_pkg(flow, environment, mfcontent):
+    """Build a bare MetaflowPackage instance with minimal state."""
     pkg = object.__new__(MetaflowPackage)
     pkg._flow = flow
     pkg._environment = environment
     pkg._mfcontent = mfcontent
+    pkg._user_content_from_addl = {}
+    return pkg
+
+
+def _call_add_addl_files(flow, environment, mfcontent):
+    """Call _add_addl_files on a bare MetaflowPackage instance."""
+    pkg = _build_pkg(flow, environment, mfcontent)
     pkg._add_addl_files()
     return mfcontent
 
@@ -204,3 +211,54 @@ def record():
         )
         _call_add_addl_files(flow, _make_environment(), _make_mfcontent())
         assert call_order == ["flow_deco", "step_deco"]
+
+
+def test_user_content_recorded():
+    """USER_CONTENT tuples are recorded in _user_content_from_addl (not sent
+    to _mfcontent which handles code/other files only)."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f:
+        deco = _make_deco([(f.name, "extra.py", ContentType.USER_CONTENT)])
+        flow = _make_flow(
+            steps=[_make_step()],
+            flow_decorators={"my_deco": [deco]},
+        )
+        mfcontent = _make_mfcontent()
+        pkg = _build_pkg(flow, _make_environment(), mfcontent)
+        pkg._add_addl_files()
+
+        # Not routed to _mfcontent — USER_CONTENT is packaged alongside user code.
+        mfcontent.add_code_file.assert_not_called()
+        mfcontent.add_other_file.assert_not_called()
+        mfcontent.add_module.assert_not_called()
+
+        assert pkg._user_content_from_addl == {"extra.py": os.path.realpath(f.name)}
+
+
+def test_user_content_duplicate_same_path_dedup():
+    """Same USER_CONTENT arcname with the same path from two decorators: dedup."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f:
+        d1 = _make_deco([(f.name, "shared.py", ContentType.USER_CONTENT)])
+        d2 = _make_deco([(f.name, "shared.py", ContentType.USER_CONTENT)])
+        flow = _make_flow(
+            steps=[_make_step()],
+            flow_decorators={"fd": [d1, d2]},
+        )
+        pkg = _build_pkg(flow, _make_environment(), _make_mfcontent())
+        pkg._add_addl_files()
+        assert pkg._user_content_from_addl == {"shared.py": os.path.realpath(f.name)}
+
+
+def test_user_content_duplicate_different_path_raises():
+    """Same USER_CONTENT arcname with different paths raises the usual exception."""
+    with tempfile.NamedTemporaryFile(suffix=".py") as f1, tempfile.NamedTemporaryFile(
+        suffix=".py"
+    ) as f2:
+        d1 = _make_deco([(f1.name, "shared.py", ContentType.USER_CONTENT)])
+        d2 = _make_deco([(f2.name, "shared.py", ContentType.USER_CONTENT)])
+        flow = _make_flow(
+            steps=[_make_step()],
+            flow_decorators={"fd": [d1, d2]},
+        )
+        with pytest.raises(NonUniqueFileNameToFilePathMappingException):
+            pkg = _build_pkg(flow, _make_environment(), _make_mfcontent())
+            pkg._add_addl_files()

test/unit/test_package_suffixes_mutator.py

@@ -0,0 +1,76 @@
+"""Tests for the ``package_suffixes`` example FlowMutator."""
+
+import os
+import tempfile
+from unittest import mock
+
+from metaflow.packaging_sys import ContentType
+from metaflow.plugins.package_suffixes_mutator import package_suffixes
+
+
+def _make_flow(tmpdir):
+    """Write a minimal flow.py and a few sibling files; return the flow file."""
+    flow_file = os.path.join(tmpdir, "flow.py")
+    with open(flow_file, "w") as f:
+        f.write("# dummy flow\n")
+    with open(os.path.join(tmpdir, "config.yaml"), "w") as f:
+        f.write("a: 1\n")
+    with open(os.path.join(tmpdir, "data.json"), "w") as f:
+        f.write("{}\n")
+    os.makedirs(os.path.join(tmpdir, "sub"))
+    with open(os.path.join(tmpdir, "sub", "nested.yaml"), "w") as f:
+        f.write("b: 2\n")
+    with open(os.path.join(tmpdir, "ignored.txt"), "w") as f:
+        f.write("ignored\n")
+    return flow_file
+
+
+def _build_mutator(suffixes, flow_file):
+    m = package_suffixes.__new__(package_suffixes)
+    # We only need _flow_cls for inspect.getfile(); patch that directly instead
+    # of constructing a real class.
+    m._flow_cls = mock.Mock()
+    m.init(suffixes)
+    return m
+
+
+def test_init_list_form():
+    m = package_suffixes.__new__(package_suffixes)
+    m.init([".yaml", "json"])
+    assert m._suffixes == (".yaml", ".json")
+
+
+def test_init_string_form():
+    m = package_suffixes.__new__(package_suffixes)
+    m.init(".yaml,json, .txt")
+    assert m._suffixes == (".yaml", ".json", ".txt")
+
+
+def test_add_to_package_yields_matching_files():
+    with tempfile.TemporaryDirectory() as tmp:
+        flow_file = _make_flow(tmp)
+        m = _build_mutator([".yaml", ".json"], flow_file)
+        with mock.patch("inspect.getfile", return_value=flow_file):
+            results = list(m.add_to_package())
+
+        # All tuples are USER_CONTENT
+        assert all(t[2] == ContentType.USER_CONTENT for t in results)
+
+        arcnames = {t[1] for t in results}
+        # walk() yields arcnames relative to the flow directory (no flow dir
+        # basename prefix), matching the convention used by _user_code_tuples.
+        assert "config.yaml" in arcnames
+        assert "data.json" in arcnames
+        assert os.path.join("sub", "nested.yaml") in arcnames
+        # Non-matching files are not included.
+        assert "ignored.txt" not in arcnames
+        # flow.py is a .py file — not part of the configured extra suffixes.
+        assert "flow.py" not in arcnames
+
+
+def test_add_to_package_empty_suffixes_yields_nothing():
+    with tempfile.TemporaryDirectory() as tmp:
+        flow_file = _make_flow(tmp)
+        m = _build_mutator([], flow_file)
+        with mock.patch("inspect.getfile", return_value=flow_file):
+            assert list(m.add_to_package()) == []