Strip fb-only comments when using opensource manifest shipit builder

Summary: Lines marked as comments for open source distribution should not be included in internal open source-esque testing.

Reviewed By: bigfootjon, itamaro

Differential Revision: D93801016

fbshipit-source-id: 344a5b0ca29b89d81a5d0ba9d961c884e28f03dc

Author: r-barnes

build/fbcode_builder/getdeps/fetcher.py

@@ -380,6 +380,31 @@ def copy_if_different(src_name, dest_name) -> bool:
     return True
 
 
+def filter_strip_marker(dest_name, marker) -> None:
+    """Strip lines/blocks tagged with the given marker from a file."""
+    try:
+        with open(dest_name, "r") as f:
+            content = f.read()
+    except (UnicodeDecodeError, PermissionError):
+        return
+
+    if marker not in content:
+        return
+
+    escaped = re.escape(marker)
+    block_re = re.compile(
+        r"[^\n]*" + escaped + r"-start[^\n]*\n.*?[^\n]*" + escaped + r"-end[^\n]*\n?",
+        re.DOTALL,
+    )
+    line_re = re.compile(r".*" + escaped + r".*\n?")
+
+    filtered = block_re.sub("", content)
+    filtered = line_re.sub("", filtered)
+    if filtered != content:
+        with open(dest_name, "w") as f:
+            f.write(filtered)
+
+
 def list_files_under_dir_newer_than_timestamp(dir_to_scan, ts):
     for root, _dirs, files in os.walk(dir_to_scan):
         for src_file in files:
@@ -394,6 +419,7 @@ def __init__(self) -> None:
         self.roots = []
         self.mapping = []
         self.exclusion = []
+        self.strip_marker = "@fb-only"
 
     def add_mapping(self, fbsource_dir, target_dir) -> None:
         """Add a posix path or pattern.  We cannot normpath the input
@@ -492,6 +518,7 @@ def st_dev(path):
                     if target_name:
                         full_file_list.add(target_name)
                         if copy_if_different(full_name, target_name):
+                            filter_strip_marker(target_name, self.strip_marker)
                             change_status.record_change(target_name)
                             if update_count < 10:
                                 print("Updated %s -> %s" % (full_name, target_name))
@@ -680,6 +707,8 @@ def update(self) -> ChangeStatus:
         for pattern in self.manifest.get_section_as_args("shipit.strip", self.ctx):
             mapping.add_exclusion(pattern)
 
+        mapping.strip_marker = self.manifest.shipit_strip_marker
+
         return mapping.mirror(self.build_options.fbsource_dir, self.repo_dir)
 
     # pyre-fixme[15]: `hash` overrides method defined in `Fetcher` inconsistently.

build/fbcode_builder/getdeps/manifest.py

@@ -51,6 +51,7 @@
             "shipit_fbcode_builder": OPTIONAL,
             "use_shipit": OPTIONAL,
             "shipit_external_branch": OPTIONAL,
+            "shipit_strip_marker": OPTIONAL,
         },
     },
     "dependencies": {"optional_section": True, "allow_values": False},
@@ -258,6 +259,9 @@ def __init__(self, file_name, fp=None):
         self.shipit_project = self.get("manifest", "shipit_project")
         self.shipit_fbcode_builder = self.get("manifest", "shipit_fbcode_builder")
         self.resolved_system_packages = {}
+        self.shipit_strip_marker = self.get(
+            "manifest", "shipit_strip_marker", defval="@fb-only"
+        )
 
         if self.name != os.path.basename(file_name):
             raise Exception(

build/fbcode_builder/getdeps/test/strip_marker_test.py

@@ -0,0 +1,159 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+
+import os
+import tempfile
+import unittest
+
+from ..fetcher import filter_strip_marker
+from ..manifest import ManifestParser
+
+
+class ManifestStripMarkerTest(unittest.TestCase):
+    def test_default_strip_marker(self) -> None:
+        p = ManifestParser(
+            "test",
+            """
+[manifest]
+name = test
+""",
+        )
+        self.assertEqual(p.shipit_strip_marker, "@fb-only")
+
+    def test_custom_strip_marker(self) -> None:
+        p = ManifestParser(
+            "test",
+            """
+[manifest]
+name = test
+shipit_strip_marker = @oss-disable
+""",
+        )
+        self.assertEqual(p.shipit_strip_marker, "@oss-disable")
+
+
+class FilterStripMarkerTest(unittest.TestCase):
+    def _write_temp(self, content: str) -> str:
+        fd, path = tempfile.mkstemp(suffix=".txt")
+        os.close(fd)
+        with open(path, "w") as f:
+            f.write(content)
+        return path
+
+    def _read(self, path: str) -> str:
+        with open(path, "r") as f:
+            return f.read()
+
+    def test_single_line_removal(self) -> None:
+        path = self._write_temp("keep this\nremove this @fb-only\nkeep this too\n")
+        try:
+            filter_strip_marker(path, "@fb-only")
+            self.assertEqual(self._read(path), "keep this\nkeep this too\n")
+        finally:
+            os.unlink(path)
+
+    def test_block_removal(self) -> None:
+        content = (
+            "before\n"
+            "// @fb-only-start\n"
+            "secret stuff\n"
+            "more secret\n"
+            "// @fb-only-end\n"
+            "after\n"
+        )
+        path = self._write_temp(content)
+        try:
+            filter_strip_marker(path, "@fb-only")
+            self.assertEqual(self._read(path), "before\nafter\n")
+        finally:
+            os.unlink(path)
+
+    def test_no_marker_present_no_change(self) -> None:
+        original = "nothing special here\njust plain code\n"
+        path = self._write_temp(original)
+        try:
+            filter_strip_marker(path, "@fb-only")
+            self.assertEqual(self._read(path), original)
+        finally:
+            os.unlink(path)
+
+    def test_custom_marker_single_line(self) -> None:
+        content = "keep\nremove @oss-disable\nkeep too\n"
+        path = self._write_temp(content)
+        try:
+            filter_strip_marker(path, "@oss-disable")
+            self.assertEqual(self._read(path), "keep\nkeep too\n")
+        finally:
+            os.unlink(path)
+
+    def test_custom_marker_block(self) -> None:
+        content = (
+            "before\n"
+            "# @oss-disable-start\n"
+            "internal only\n"
+            "# @oss-disable-end\n"
+            "after\n"
+        )
+        path = self._write_temp(content)
+        try:
+            filter_strip_marker(path, "@oss-disable")
+            self.assertEqual(self._read(path), "before\nafter\n")
+        finally:
+            os.unlink(path)
+
+    def test_custom_marker_ignores_default(self) -> None:
+        """When using a custom marker, @fb-only lines should be kept."""
+        content = "keep @fb-only\nremove @oss-disable\nplain\n"
+        path = self._write_temp(content)
+        try:
+            filter_strip_marker(path, "@oss-disable")
+            self.assertEqual(self._read(path), "keep @fb-only\nplain\n")
+        finally:
+            os.unlink(path)
+
+    def test_mixed_single_and_block(self) -> None:
+        content = (
+            "line1\n"
+            "line2 @fb-only\n"
+            "line3\n"
+            "// @fb-only-start\n"
+            "block content\n"
+            "// @fb-only-end\n"
+            "line4\n"
+        )
+        path = self._write_temp(content)
+        try:
+            filter_strip_marker(path, "@fb-only")
+            self.assertEqual(self._read(path), "line1\nline3\nline4\n")
+        finally:
+            os.unlink(path)
+
+    def test_marker_with_regex_metacharacters(self) -> None:
+        """Markers containing regex metacharacters should be escaped properly."""
+        content = "keep\nremove @fb.only\nkeep too\n"
+        path = self._write_temp(content)
+        try:
+            # With proper escaping, the dot is literal, not a wildcard
+            filter_strip_marker(path, "@fb.only")
+            self.assertEqual(self._read(path), "keep\nkeep too\n")
+        finally:
+            os.unlink(path)
+
+    def test_binary_file_skipped(self) -> None:
+        """Binary files that can't be decoded as UTF-8 should be skipped."""
+        fd, path = tempfile.mkstemp(suffix=".bin")
+        os.close(fd)
+        binary_content = b"\x80\x81\x82\xff\xfe"
+        with open(path, "wb") as f:
+            f.write(binary_content)
+        try:
+            filter_strip_marker(path, "@fb-only")
+            with open(path, "rb") as f:
+                self.assertEqual(f.read(), binary_content)
+        finally:
+            os.unlink(path)