feat(pre-commit): add internal-import alias check with auto-fix by dengqiaoyu · Pull Request #49 · apple/coreai-optimization

dengqiaoyu · 2026-07-20T19:16:27Z

Add an AST-based hook enforcing code style guide §3.3: public modules must _-alias symbols imported from private modules; private modules must not
Auto-fix offending import statements and rename in-file references, refusing the fix only when the bound name is shadowed
Register the hook in .pre-commit-config.yaml and document it under §3.3 "Enforcement"
Apply across src/coreai_opt/, normalizing 27 modules with missing or redundant _ aliases

crowbat · 2026-07-21T16:19:04Z

+            yield from ast.walk(stmt)
+
+
+def _locally_binds_name(node: ast.AST, name: str) -> bool:


2 comments on this function:

It looks like it misses out on the case where t in node.targets itself is a tuple/list and when node is a comprehension. Updating the checks to use _is_named_target could dedupe some logic and fix these issues at the same time:

def _locally_binds_name(node: ast.AST, name: str) -> bool: if isinstance(node, ast.Assign): return any(_is_named_target(t, name) for t in node.targets) if isinstance(node, (ast.AnnAssign, ast.AugAssign, ast.NamedExpr, ast.For, ast.AsyncFor)): return _is_named_target(node.target, name) if isinstance(node, (ast.With, ast.AsyncWith)): return any( item.optional_vars is not None and _is_named_target(item.optional_vars, name) for item in node.items ) if isinstance(node, ast.comprehension): return _is_named_target(node.target, name) return False

This function and _binds_assignment_target share very similar logic, could they be combined into one?

It looks like it misses out on the case

Good catch. I think I was missing a few cases here:

a, helper = f() for a, helper in items with cm as (a, helper) [y for helper in items]

This function and _binds_assignment_target share very similar logic, could they be combined into one?

Before your previous comment, these two had one difference. After adding the cases above, they can now be merged into one. I'll make that change.

crowbat · 2026-07-21T16:38:20Z

+
+def _node_span(node: _Positioned, line_starts: list[int]) -> tuple[int, int]:
+    """Convert an AST node's (lineno, col_offset, end_lineno, end_col_offset) to byte offsets."""
+    start = line_starts[node.lineno - 1] + node.col_offset


Somewhat a corner case, but if the code has some character that is non-ASCII that uses more than one byte, the logic here can result in a malformed update:

# Before the fix from temp_pkg._priv import helper def use(): café = ... return café + helper() # After the fix from temp_pkg._priv import helper as _helper def use(): café = ... return café + h_helper)

I doubt people will use non-ASCII characters when writing code, but we should still handle this case. Will change it.

crowbat · 2026-07-21T16:41:07Z

+    """
+    violations: list[Violation] = []
+
+    for node in ast.walk(tree):


Should we also check for and flag plain imports of private packages? Like import pkg._private

Looks like we are missing checks for relative imports (from ._priv import x)

crowbat · 2026-07-21T16:45:13Z

+    """
+    original = alias.name
+    current_alias = alias.asname
+    if original.startswith("_") or current_alias not in (None, original, "_" + original):


Another corner case, this currently does not handle the case if someone names things with a custom alias (from a._b import c as d)

crowbat · 2026-07-21T16:47:13Z

+    )
+
+
+def _statement_edits(


If an import line has a comment at the end, looks like the comment would be removed as part of the rename

crowbat · 2026-07-21T16:50:49Z

+          symbols imported from private modules with a `_` prefix; private
+          modules must not. Auto-fixes the import statement and renames
+          references in the same file.
+        entry: python scripts/pre_commit/check_internal_import_aliases.py --fix


In general it seems like there are a lot of potential corner cases associated with automatically fixing the code. What do you think about making --fix opt-in instead where the user manually runs it if they want to have things be automatically fixed?

We could have the pre-commit hook only do the detection part of it to fail the commit if it finds any transgressions. I'm thinking that if any automatic fixes are made which break the code, other pre-commit hooks may also have run (linting, line too long, etc.) which would conflate the changes if the user needs to revert the code back to before this hook took action.

What do you think about making --fix opt-in

Agreed. There are many corner cases that this code cannot comprehensively cover, it's essentially acting as an AST parser, and I don't want it to be too complicated .

~~I'll make the pre-commit hook a warning instead of an error. This way, it won't fail the commit,~~ but it will generate error information with correction suggestions. I'll also add a --fix flag to let users fix everything in one go.

Thanks for the suggestion.

Here's what I came up with:

For the two attached files, _alias_demo.py and alias_demo.py, here's what we'd get if we committed them directly with the pre-commit hook added:

Check `_`-prefix aliasing of internal imports..............................Failed - hook id: check-internal-import-aliases - exit code: 1 src/coreai_opt/_alias_demo.py:14: Unnecessary `_` alias for `CompressionType` imported from private module `coreai_opt._utils.metadata_utils` into a private module.; Unnecessary `_` alias for `MILCompressionMetadata` imported from private module `coreai_opt._utils.metadata_utils` into a private module. src/coreai_opt/_alias_demo.py:18: Unnecessary `_` alias for `move_model_to_eval` imported from private module `coreai_opt._utils.torch_utils` into a private module. src/coreai_opt/_alias_demo.py:19: Unnecessary `_` alias for `version_ge` imported from private module `coreai_opt._utils.version_utils` into a private module. src/coreai_opt/_alias_demo.py:21: Unnecessary `_` alias for `PartialConstructor` imported from private module `._utils.spec_utils` into a private module. src/coreai_opt/alias_demo.py:16: Missing `_` alias for private module `coreai_opt._utils.registry_utils` bound as `registry` in a public module. Use `import coreai_opt._utils.registry_utils as _registry`. src/coreai_opt/alias_demo.py:17: Missing `_` alias for `CompressionType` imported from private module `coreai_opt._utils.metadata_utils` into a public module. Use `CompressionType as _CompressionType`.; Missing `_` alias for `MILCompressionMetadata` imported from private module `coreai_opt._utils.metadata_utils` into a public module. Use `MILCompressionMetadata as _MILCompressionMetadata`. src/coreai_opt/alias_demo.py:21: Missing `_` alias for `PartialConstructor` imported from private module `coreai_opt._utils.spec_utils` into a public module. Use `PartialConstructor as _PC`. src/coreai_opt/alias_demo.py:22: Missing `_` alias for `move_model_to_eval` imported from private module `coreai_opt._utils.torch_utils` into a public module. Use `move_model_to_eval as _move_model_to_eval`. (needs a manual fix: bound name is shadowed) src/coreai_opt/alias_demo.py:23: Missing `_` alias for `version_ge` imported from private module `coreai_opt._utils.version_utils` into a public module. Use `version_ge as _version_ge`. src/coreai_opt/alias_demo.py:25: Missing `_` alias for `lazy_import_coreai_torch` imported from private module `._utils.import_utils` into a public module. Use `lazy_import_coreai_torch as _lazy_import_coreai_torch`. 10 issue(s) found. Fix them before committing — run `python scripts/pre_commit/check_internal_import_aliases.py --fix` to apply the fixable ones automatically.

And here's the output from running the auto-fix:

$ python scripts/pre_commit/check_internal_import_aliases.py --fix src/coreai_opt/_alias_demo.py:14: FIXED — Unnecessary `_` alias for `CompressionType` imported from private module `coreai_opt._utils.metadata_utils` into a private module.; Unnecessary `_` alias for `MILCompressionMetadata` imported from private module `coreai_opt._utils.metadata_utils` into a private module. src/coreai_opt/_alias_demo.py:18: FIXED — Unnecessary `_` alias for `move_model_to_eval` imported from private module `coreai_opt._utils.torch_utils` into a private module. src/coreai_opt/_alias_demo.py:19: FIXED — Unnecessary `_` alias for `version_ge` imported from private module `coreai_opt._utils.version_utils` into a private module. src/coreai_opt/_alias_demo.py:21: FIXED — Unnecessary `_` alias for `PartialConstructor` imported from private module `._utils.spec_utils` into a private module. src/coreai_opt/alias_demo.py:16: FIXED — Missing `_` alias for private module `coreai_opt._utils.registry_utils` bound as `registry` in a public module. Use `import coreai_opt._utils.registry_utils as _registry`. src/coreai_opt/alias_demo.py:17: FIXED — Missing `_` alias for `CompressionType` imported from private module `coreai_opt._utils.metadata_utils` into a public module. Use `CompressionType as _CompressionType`.; Missing `_` alias for `MILCompressionMetadata` imported from private module `coreai_opt._utils.metadata_utils` into a public module. Use `MILCompressionMetadata as _MILCompressionMetadata`. src/coreai_opt/alias_demo.py:21: FIXED — Missing `_` alias for `PartialConstructor` imported from private module `coreai_opt._utils.spec_utils` into a public module. Use `PartialConstructor as _PC`. src/coreai_opt/alias_demo.py:23: FIXED — Missing `_` alias for `version_ge` imported from private module `coreai_opt._utils.version_utils` into a public module. Use `version_ge as _version_ge`. src/coreai_opt/alias_demo.py:25: FIXED — Missing `_` alias for `lazy_import_coreai_torch` imported from private module `._utils.import_utils` into a public module. Use `lazy_import_coreai_torch as _lazy_import_coreai_torch`. src/coreai_opt/alias_demo.py:22: Missing `_` alias for `move_model_to_eval` imported from private module `coreai_opt._utils.torch_utils` into a public module. Use `move_model_to_eval as _move_model_to_eval`. (needs a manual fix: bound name is shadowed) 1 issue(s) could not be auto-fixed; resolve them by hand.

_alias_demo.py
alias_demo.py

Makes sense, I do like keeping it as an error instead of just a warning

- Add an AST-based hook enforcing code style guide §3.3: public modules must alias symbols imported from a private module with a `_` prefix; private modules must not - On commit it fails and reports each violation with a suggested fix, but does not edit files — run the script with `--fix` to apply fixes manually (kept opt-in so its rewrites don't get entangled with the formatter/linter hooks) - Detect absolute, relative, and aliased plain imports, plus custom aliases; preserve comments and handle multibyte source via byte offsets; skip cases that are unsafe to rename (shadowed or rebound names) - Add a (before, after) test suite under `tests/devtools/` covering every case

- Add the required `_` alias to private-module imports in public modules, drop the redundant `_` alias in private modules, and rename references to match, per code style guide §3.3 - Generated with `python scripts/pre_commit/check_internal_import_aliases.py --fix` - Kept as a separate commit from the hook so a future rebase conflict here can be resolved by reverting `src/` and re-running `--fix`, instead of hand-merging auto-generated edits

crowbat

Looks good to me, thanks for the fixes!

dengqiaoyu requested review from crowbat and guru-desh July 20, 2026 19:16

dengqiaoyu assigned crowbat Jul 20, 2026

dengqiaoyu added documentation Improvements or additions to documentation enhancement New feature or request labels Jul 20, 2026

crowbat reviewed Jul 21, 2026

View reviewed changes

dengqiaoyu added 2 commits July 22, 2026 17:19

dengqiaoyu force-pushed the u/qiaoyu_deng/private-import-check branch from 23e77b7 to 7c01c37 Compare July 23, 2026 00:22

crowbat approved these changes Jul 23, 2026

View reviewed changes

dengqiaoyu merged commit 2d92c5c into apple:main Jul 23, 2026
11 checks passed

dengqiaoyu deleted the u/qiaoyu_deng/private-import-check branch July 23, 2026 17:12

		yield from ast.walk(stmt)


		def _locally_binds_name(node: ast.AST, name: str) -> bool:

Uh oh!

Conversation

dengqiaoyu commented Jul 20, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dengqiaoyu Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dengqiaoyu Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

crowbat left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dengqiaoyu Jul 22, 2026 •

edited

Loading

dengqiaoyu Jul 22, 2026 •

edited

Loading