diff --git a/packages/syft-restrict/src/syft_restrict/astutil.py b/packages/syft-restrict/src/syft_restrict/astutil.py index 7923db9f02b..3e0ca3e790b 100644 --- a/packages/syft-restrict/src/syft_restrict/astutil.py +++ b/packages/syft-restrict/src/syft_restrict/astutil.py @@ -138,6 +138,9 @@ def scan_file(tree: ast.Module, private_ranges) -> FileScan: for node in ast.walk(tree): if isinstance(node, ast.Import): for alias in node.names: + # The alias (key) is the name the private code writes; the value is the + # fully-qualified path the policy matches against. _Checker._resolve() swaps a + # call's root name for this value and keeps the rest of the chain verbatim. # Bind the name Python actually binds at runtime: # `import jax.numpy as jnp` binds `jnp` -> the jax.numpy module # `import jax.numpy` binds `jax` -> the jax PACKAGE (not jax.numpy!) -- diff --git a/packages/syft-restrict/src/syft_restrict/verifier.py b/packages/syft-restrict/src/syft_restrict/verifier.py index 4ed5a07b560..e10c2d1a254 100644 --- a/packages/syft-restrict/src/syft_restrict/verifier.py +++ b/packages/syft-restrict/src/syft_restrict/verifier.py @@ -963,8 +963,20 @@ def _check_reserved_name(self, node: ast.AST, name: str) -> None: # ── path resolution ────────────────────────────────────────────────────────────────────── def _resolve(self, path: str) -> str: - """Rewrite a dotted path's import alias to its fully-qualified form - (`jnp.einsum` -> `jax.numpy.einsum`).""" + """Rewrite a call's dotted path to the fully-qualified form the policy matches against. + + Two pieces of state meet here. The call's own dotted text is what the private code *writes* + (`jnp.einsum`, `jax.numpy.save`). ``import_bindings`` (built in ``scan_file``) maps each + imported *name* -> the fully-qualified path it stands for. We split the call into its root + name + the rest of the chain, swap the root for its binding, and keep the rest verbatim: + + `import jax.numpy as jnp` binds jnp -> "jax.numpy"; `jnp.einsum` -> `jax.numpy.einsum` + `import jax.numpy` binds jax -> "jax"; `jax.numpy.save` -> `jax.numpy.save` + + Only the alias form rewrites the name to a longer path; a plain `import a.b` binds the root + to itself because the call already spells the full path, so we just pass `rest` through. The + returned path is what allow_functions/disallow_functions are then checked against. + """ root, _, rest = path.partition(".") base = self.scan.import_bindings.get(root, root) return f"{base}.{rest}" if rest else base