fix: keep obfuscated copy row-aligned with the source

packages/syft-restrict/src/syft_restrict/__init__.py

@@ -10,7 +10,13 @@
 __version__ = "0.1.0"
 
 from .auditor import AuditReport, PathAudit, audit_allow_functions
-from .errors import MarkerError, PolicyViolation, RestrictError, SourceEncodingError
+from .errors import (
+    MarkerError,
+    ObfuscationError,
+    PolicyViolation,
+    RestrictError,
+    SourceEncodingError,
+)
 from .markers import parse_markers
 from .obfuscator import obfuscate
 from .policy import Policy
@@ -32,6 +38,7 @@
     "PolicyViolation",
     "MarkerError",
     "SourceEncodingError",
+    "ObfuscationError",
     "RestrictError",
     "__version__",
 ]

packages/syft-restrict/src/syft_restrict/astutil.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import ast
+import io
 
 from pydantic import BaseModel
 
@@ -22,6 +23,15 @@ def normalize_ranges(private) -> list[tuple[int, int]]:
     return ranges
 
 
+def source_rows(text: str) -> list[str]:
+    """Split ``text`` into rows on ``\\n`` only, as ``tokenize`` counts them.
+
+    Each row keeps its ``\\n``. Unlike ``str.splitlines()``, a form feed or another separator
+    character does not end a row.
+    """
+    return io.StringIO(text).readlines()
+
+
 def row_in_ranges(row: int, ranges) -> bool:
     """True if a 1-based line number falls inside any range."""
     return any(lo <= row <= hi for lo, hi in ranges)

packages/syft-restrict/src/syft_restrict/errors.py

@@ -30,3 +30,7 @@ class MarkerError(RestrictError):
 
 class SourceEncodingError(RestrictError):
     """Raised when a source file is not plain UTF-8."""
+
+
+class ObfuscationError(RestrictError):
+    """Raised when the display copy does not keep the rows of the source."""

packages/syft-restrict/src/syft_restrict/markers.py

@@ -24,6 +24,7 @@ def attention(x):
 import re
 import tokenize
 
+from .astutil import source_rows
 from .errors import MarkerError
 
 __all__ = ["parse_markers"]
@@ -116,7 +117,7 @@ def __init__(self) -> None:
 
     def feed(self, source: str) -> None:
         markers = _scan_markers(source)
-        for line in range(1, len(source.splitlines()) + 1):
+        for line in range(1, len(source_rows(source)) + 1):
             marker = markers.get(line)
             if marker is None:
                 if self.stack:

packages/syft-restrict/src/syft_restrict/obfuscator.py

@@ -23,6 +23,7 @@
     node_in_ranges,
     normalize_ranges,
     row_in_ranges,
+    source_rows,
 )
 from .policy import DEFAULT_KEEP
 
@@ -112,13 +113,16 @@ def obfuscate(source: str, obfuscate_ranges, hide_ranges, scan: FileScan) -> str
             if new is not None:
                 edits.append((srow, scol, erow, ecol, new))
         elif tok.type == tokenize.STRING:
-            edits.append((srow, scol, erow, ecol, f'"{_BLANK}"'))
+            # A string that spans rows keeps all of them, so every later row stays where it was.
+            gap = "\n" * (erow - srow)
+            new = f'"""{_BLANK}{gap}"""' if gap else f'"{_BLANK}"'
+            edits.append((srow, scol, erow, ecol, new))
         elif tok.type == _FSTRING_MIDDLE:
             # The f-string's literal text between `{...}` parts; FSTRING_START/END (the f"/" quote
             # delimiters) are left as-is, and any interpolated expression's own tokens (NAME, etc.)
             # still flow through the normal token handling above/below like any other reference.
             if tok.string:
-                edits.append((srow, scol, erow, ecol, _BLANK))
+                edits.append((srow, scol, erow, ecol, _BLANK + "\n" * (erow - srow)))
         elif tok.type == tokenize.NUMBER:
             edits.append((srow, scol, erow, ecol, _BLANK))
         elif tok.type == tokenize.COMMENT:
@@ -225,7 +229,7 @@ def _names_of(tree, node_type, ranges) -> set[str]:
 
 # ── apply position edits to the source, preserving non-private lines verbatim ─────────────
 def _apply_edits(source: str, edits) -> str:
-    lines = source.splitlines(keepends=True)
+    lines = source_rows(source)
     # apply bottom-up so earlier edits don't shift later line indices / columns
     for srow, scol, erow, ecol, new in sorted(edits, reverse=True):
         if srow == erow:
@@ -239,7 +243,7 @@ def _apply_edits(source: str, edits) -> str:
 
 # ── blank whole hidden lines, keeping indentation; leave blank lines untouched ────────────
 def _apply_hides(text: str, hide_ranges) -> str:
-    lines = text.splitlines(keepends=True)
+    lines = source_rows(text)
     in_block = False  # inside a run of consecutive hidden line numbers
     noted = False  # has the explanatory note been added for the current block yet
     for i, line in enumerate(lines, 1):

packages/syft-restrict/src/syft_restrict/runner.py

@@ -11,9 +11,9 @@
 
 from pydantic import BaseModel, Field
 
-from .astutil import normalize_ranges, scan_file
+from .astutil import normalize_ranges, scan_file, source_rows
 from .auditor import AuditReport, audit_allow_functions
-from .errors import PolicyViolation, SourceEncodingError
+from .errors import ObfuscationError, PolicyViolation, SourceEncodingError
 from .markers import parse_markers
 from .obfuscator import obfuscate as _obfuscate
 from .policy import Policy
@@ -143,6 +143,7 @@ def _run(
 
     scan = scan_file(ast.parse(source), normalize_ranges(private))
     obfuscated = _obfuscate(source, obfuscate_ranges, hide_ranges, scan)
+    _require_same_row_count(source, obfuscated)
 
     out_path = Path(out) if out is not None else path.with_suffix(".obfuscated.py")
     out_path.write_text(obfuscated, encoding="utf-8")
@@ -165,6 +166,17 @@ def _run(
     )
 
 
+def _require_same_row_count(source: str, obfuscated: str) -> None:
+    """Raise ``ObfuscationError`` if ``obfuscated`` and ``source`` have different row counts."""
+    # The hide ranges and the certificate's ranges count rows of the source. A copy with a
+    # different row count puts the hiding on the wrong rows.
+    n_src, n_out = len(source_rows(source)), len(source_rows(obfuscated))
+    if n_out != n_src:
+        raise ObfuscationError(
+            f"the obfuscated copy has {n_out} rows and the source has {n_src}"
+        )
+
+
 def _read_source(path: Path) -> str:
     """Read the file at ``path`` and return its text, with every line ending turned into ``\\n``.
 

packages/syft-restrict/tests/obfuscate/test_obfuscate.py

@@ -4,8 +4,9 @@
 import shutil
 from pathlib import Path
 
+import pytest
 from syft_restrict import obfuscate
-from syft_restrict.astutil import scan_file
+from syft_restrict.astutil import row_in_ranges, scan_file
 from syft_restrict.runner import _run
 
 FIXTURES = Path(__file__).parents[1] / "fixtures"
@@ -103,6 +104,67 @@ def test_fstring_literal_text_is_blanked():
     assert "■" in obf
 
 
+def _assert_rows_aligned(source: str, obf: str, obfuscate_ranges, hide_ranges):
+    """Assert that public rows print unchanged and hidden code rows print as the hide marker."""
+    src_rows, obf_rows = source.split("\n"), obf.split("\n")
+    assert len(obf_rows) == len(src_rows), obf
+    for row, (src_row, obf_row) in enumerate(zip(src_rows, obf_rows), 1):
+        if row_in_ranges(row, hide_ranges) and src_row.strip():
+            assert obf_row.lstrip().startswith("■■■■■■■■"), (row, obf)
+        elif not row_in_ranges(row, [*obfuscate_ranges, *hide_ranges]):
+            assert obf_row == src_row, (row, obf)
+
+
+@pytest.mark.parametrize(
+    "obfuscate_ranges, hide_ranges",
+    [([[3, 6]], []), ([[3, 3]], [[4, 6]])],
+    ids=["obfuscate", "obfuscate-then-hide"],
+)
+def test_formfeed_keeps_row_count(obfuscate_ranges, hide_ranges):
+    # A form feed must stay part of its row, the way tokenize counts rows.
+    source = (
+        "import jax.numpy as jnp\n"
+        "# page\x0cbreak\n"
+        "def scale(x):\n"
+        "    y = x * 2\n"
+        "    z = y + 1\n"
+        "    return z\n"
+        "\n"
+        "def main(x):\n"
+        "    return scale(x)\n"
+    )
+    private = [*obfuscate_ranges, *hide_ranges]
+    scan = scan_file(ast.parse(source), private)
+    obf = obfuscate(source, obfuscate_ranges, hide_ranges, scan)
+    _assert_rows_aligned(source, obf, obfuscate_ranges, hide_ranges)
+
+
+@pytest.mark.parametrize(
+    "obfuscate_ranges, hide_ranges",
+    [([[3, 6]], []), ([[3, 6]], [[7, 7]])],
+    ids=["obfuscate", "obfuscate-then-hide"],
+)
+def test_multiline_docstring_keeps_row_count(obfuscate_ranges, hide_ranges):
+    # A string that spans several rows must keep all of them when it is blanked.
+    source = (
+        "import jax.numpy as jnp\n"
+        "\n"
+        "class Model:\n"
+        '    """Line one.\n'
+        "    Line two.\n"
+        '    Line three."""\n'
+        "SECRET = 1\n"
+        'PUBLIC = "public"\n'
+        "\n"
+        "def main():\n"
+        "    return PUBLIC\n"
+    )
+    private = [*obfuscate_ranges, *hide_ranges]
+    scan = scan_file(ast.parse(source), private)
+    obf = obfuscate(source, obfuscate_ranges, hide_ranges, scan)
+    _assert_rows_aligned(source, obf, obfuscate_ranges, hide_ranges)
+
+
 def test_obfuscation_is_deterministic(tmp_path):
     dir_a = tmp_path / "a"
     dir_b = tmp_path / "b"

packages/syft-restrict/tests/test_run.py

@@ -6,7 +6,14 @@
 from pathlib import Path
 
 import pytest
-from syft_restrict import MarkerError, PolicyViolation, SourceEncodingError, run
+from syft_restrict import (
+    MarkerError,
+    ObfuscationError,
+    PolicyViolation,
+    SourceEncodingError,
+    obfuscate,
+    run,
+)
 from syft_restrict.runner import _run
 from verify.helpers import normalize_source
 
@@ -195,6 +202,26 @@ def test__run_uses_explicit_ranges_and_ignores_markers(tmp_path):
     assert result.ok
 
 
+def test_run_refuses_copy_with_different_row_count(tmp_path, monkeypatch):
+    # No known input makes obfuscate() change the row count, so the test drops a row itself.
+    src = tmp_path / "model.py"
+    shutil.copy(FIXTURES / "marked_model.py", src)
+    monkeypatch.setattr(
+        "syft_restrict.runner._obfuscate",
+        lambda *args: obfuscate(*args).split("\n", 1)[1],
+    )
+
+    # A broken copy is an obfuscator fault, not a verification result, so strict=False raises too.
+    with pytest.raises(ObfuscationError):
+        run(
+            src,
+            allow_functions=ALLOW_FUNCTIONS,
+            allow_operators=ALLOW_OPERATORS,
+            strict=False,
+        )
+    assert not (tmp_path / "model.obfuscated.py").exists()
+
+
 def test_run_refuses_coding_cookie_that_changes_program(tmp_path):
     # The enclave loads model code with importlib, which honours a PEP 263 coding cookie.
     src = tmp_path / "model.py"