tools/fork/strip.py

@@ -239,6 +239,26 @@ def deactivate_enterprise_in_scripts(tree):
     return edits
 
 
+def trailing_attribute(lines):
+    """
+    How many lines at the end of `lines` make up one outer attribute, or 0.
+
+    An attribute can span lines, e.g. a `#[cfg(all(` whose conditions follow
+    one per line. Leaving its tail behind would hand the attribute to the next
+    item, and compile that item out instead of the removed `mod`.
+    """
+    if not lines or not lines[-1].rstrip().endswith(']'):
+        return 0
+    for i in range(len(lines) - 1, -1, -1):
+        text = lines[i].strip()
+        if not text or text.endswith((';', '{', '}')):
+            return 0
+        if text.startswith('#['):
+            joined = ''.join(l.strip() for l in lines[i:])
+            return len(lines) - i if joined.count('[') == joined.count(']') else 0
+    return 0
+
+
 def remove_dangling_mods(tree):
     """
     Remove `mod x;` declarations whose file ossify.py deleted.
@@ -259,8 +279,9 @@ def remove_dangling_mods(tree):
             m = decl.match(line)
             if m and not (base / f'{m.group(1)}.rs').is_file() and not (base / m.group(1) / 'mod.rs').is_file():
                 dropped = [line]
-                while keep and keep[-1].strip().startswith('#['):
-                    dropped.insert(0, keep.pop())
+                while (n_attr := trailing_attribute(keep)):
+                    dropped[:0] = keep[-n_attr:]
+                    del keep[-n_attr:]
                 removed.append({'file': str(path.relative_to(tree)), 'line': n, 'module': m.group(1),
                                 'lines': [l.strip() for l in dropped]})
                 continue

tools/fork/tests/test_strip.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: 2026 Coffey Labs LLC
+# SPDX-License-Identifier: AGPL-3.0-only
+"""Tests for tools/fork/strip.py: python3 -m unittest discover tools/fork/tests"""
+import importlib.util
+import os
+import pathlib
+import tempfile
+import unittest
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+spec = importlib.util.spec_from_file_location('strip', os.path.join(HERE, '..', 'strip.py'))
+strip = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(strip)
+
+MULTILINE = """use std::sync::Arc;
+
+#[cfg(all(
+    test,
+    feature = "postgres",
+))]
+#[rustfmt::skip]
+#[path = "../../missing/contention_tests.rs"]
+mod contention_tests;
+
+pub struct Batch {
+    pub ops: Vec<u8>,
+}
+"""
+
+
+class RemoveDanglingMods(unittest.TestCase):
+    def strip_file(self, text):
+        with tempfile.TemporaryDirectory() as d:
+            src = pathlib.Path(d, 'src')
+            src.mkdir()
+            path = src / 'lib.rs'
+            path.write_text(text, encoding='utf-8')
+            removed = strip.remove_dangling_mods(pathlib.Path(d))
+            return path.read_text(encoding='utf-8'), removed
+
+    def test_multiline_attribute_goes_with_its_mod(self):
+        out, removed = self.strip_file(MULTILINE)
+        self.assertNotIn('cfg(all(', out)
+        self.assertNotIn('feature = "postgres"', out)
+        self.assertIn('pub struct Batch {', out)
+        self.assertEqual(len(removed), 1)
+        self.assertEqual(removed[0]['lines'][0], '#[cfg(all(')
+
+    def test_single_line_attributes_still_removed(self):
+        out, _ = self.strip_file('#[cfg(test)]\nmod tests;\n\nfn kept() {}\n')
+        self.assertEqual(out, '\nfn kept() {}\n')
+
+    def test_previous_item_is_untouched(self):
+        out, _ = self.strip_file('const LIMIT: [u8; 2] = [0, 1];\nmod gone;\n')
+        self.assertEqual(out, 'const LIMIT: [u8; 2] = [0, 1];\n')
+
+    def test_existing_module_kept(self):
+        with tempfile.TemporaryDirectory() as d:
+            src = pathlib.Path(d, 'src')
+            src.mkdir()
+            (src / 'here.rs').write_text('', encoding='utf-8')
+            (src / 'lib.rs').write_text('#[cfg(all(\n    test,\n))]\nmod here;\n', encoding='utf-8')
+            self.assertEqual(strip.remove_dangling_mods(pathlib.Path(d)), [])
+
+
+if __name__ == '__main__':
+    unittest.main()