From 71c9286ef637f5e36aeeed1044ce4f87b75f44cb Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 25 Sep 2026 23:49:32 +0000 Subject: [PATCH] Don't drop comments inside __all__, and only read markers from comments - Rewriting a multi-line __all__ rebuilt the literal and deleted comments inside it. Such an __all__ is now reported but not rewritten. - The message for an __all__ that can't be rewritten said it was built from several statements, also when it was a non-literal like ["A"] + ["B"]; it now names all the reasons. - `# unexport: not-public` / `public` were matched against whole physical lines, so the text inside a string literal marked the name. Only comment tokens count now. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_019P9bvuGwAyVCNUdsAYPB1V --- docs/changelog.md | 2 ++ src/unexport/analyzer.py | 14 +++++++------- src/unexport/main.py | 3 ++- src/unexport/refactor.py | 16 ++++++++++++++++ tests/test_analyzer.py | 7 +++++++ tests/test_refactor.py | 34 ++++++++++++++++++++++++++++++++++ 6 files changed, 68 insertions(+), 8 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index cf5461c..a7041eb 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -54,6 +54,8 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm reported but not rewritten, and a dynamic `__all__` (e.g. `["a"] + sub.__all__`) is left alone [#41](https://github.com/hakancelikdev/unexport/issues/41) - Crash on modules whose `__all__` is a set +- An `__all__` with comments inside the literal is reported but not rewritten, so the + comments are not lost; `# unexport:` markers inside string literals are ignored ## [0.4.0] - 2022-11-05 diff --git a/src/unexport/analyzer.py b/src/unexport/analyzer.py index f0ebfb7..c7177a8 100644 --- a/src/unexport/analyzer.py +++ b/src/unexport/analyzer.py @@ -173,13 +173,13 @@ def is_dynamic_all(self) -> bool: def set_extra_attr(self, tree: ast.AST) -> None: skip, add = set(), set() readline = io.StringIO(self.source).readline - for _, _, start, _, line in tokenize.generate_tokens(readline): - if re.search(C.SKIP_COMMENTS_REGEX_PATTERN, line, re.IGNORECASE): - lineno = start[0] - skip.add(lineno) - if re.search(C.ADD_COMMENTS_REGEX_PATTERN, line, re.IGNORECASE): - lineno = start[0] - add.add(lineno) + for token in tokenize.generate_tokens(readline): + if token.type != tokenize.COMMENT: # not the same text inside a string literal + continue + if re.search(C.SKIP_COMMENTS_REGEX_PATTERN, token.string, re.IGNORECASE): + skip.add(token.start[0]) + if re.search(C.ADD_COMMENTS_REGEX_PATTERN, token.string, re.IGNORECASE): + add.add(token.start[0]) for node in ast.walk(tree): if isinstance(node, C.ALL_NODE) and node.lineno in skip: diff --git a/src/unexport/main.py b/src/unexport/main.py index 3cc1f39..c2492e2 100644 --- a/src/unexport/main.py +++ b/src/unexport/main.py @@ -86,7 +86,8 @@ def main(argv: Sequence[str] | None = None) -> int: if new_source == source: print( color.paint(py_path.as_posix(), color.YELLOW) - + ": __all__ is built from several statements and can't be updated automatically; expected " + + ": __all__ can't be updated automatically (it is built from several statements, is not a" + + " plain list, tuple or set, or has comments inside); expected " + color.paint("__all__ = " + str(expected_all), color.GREEN) ) continue diff --git a/src/unexport/refactor.py b/src/unexport/refactor.py index 2c5f929..898be66 100644 --- a/src/unexport/refactor.py +++ b/src/unexport/refactor.py @@ -1,7 +1,9 @@ from __future__ import annotations import ast +import io import re +import tokenize from unexport.dunder_all import find_all_statements @@ -64,6 +66,18 @@ def _replace_value(lines: list[str], node: ast.expr, expected_all: list[str]) -> lines[start : end + 1] = [prefix + text + suffix] +def _has_comments(source: str, lines: list[str], node: ast.expr) -> bool: + """Whether there are comments inside the (multi-line) literal.""" + start, end = node.lineno, node.end_lineno or node.lineno + # AST column offsets are UTF-8 byte offsets, tokenize's are character offsets. + end_column = len(lines[end - 1].encode()[: node.end_col_offset].decode()) + for token in tokenize.generate_tokens(io.StringIO(source).readline): + if token.type == tokenize.COMMENT and start <= token.start[0] <= end: + if token.start[0] < end or token.start[1] < end_column: + return True + return False + + def refactor_source(source: str, expected_all: list[str]) -> str: tree = ast.parse(source) lines = ast._splitlines_no_ff(source) # type: ignore @@ -75,6 +89,8 @@ def refactor_source(source: str, expected_all: list[str]) -> str: return source if not isinstance(node.value, (ast.List, ast.Tuple, ast.Set)): return source # e.g. ["a"] + sub.__all__ + if _has_comments(source, lines, node.value): + return source # rewriting the literal would drop them # Also when nothing is public anymore: a stale __all__ becomes empty. _replace_value(lines, node.value, expected_all) return "".join(lines) diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index 7244f7c..22835c3 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -342,6 +342,13 @@ def test_listed_submodules_of_a_package_are_kept(self): analyzer.traverse() self.assertListEqual(analyzer.expected_all, ["Thing", "sub", "subpackage"]) + def test_marker_inside_a_string_is_not_a_comment(self): + source = """\ + Message = "# unexport: not-public" + Other = 1 + """ + self.assertListEqual(self.expected_all(source), ["Message", "Other"]) + def test_listed_but_undefined_is_removed(self): source = """\ __all__ = ["gone", "func"] diff --git a/tests/test_refactor.py b/tests/test_refactor.py index a846cf9..e31a277 100644 --- a/tests/test_refactor.py +++ b/tests/test_refactor.py @@ -509,6 +509,40 @@ class Extra: ... class Extra: ... """, ), + ( # comments inside the literal would be lost, so it is left alone + """\ + __all__ = [ + "A", # the main entry point + # "Old", + ] + + class A: ... + class B: ... + """, + """\ + __all__ = [ + "A", # the main entry point + # "Old", + ] + + class A: ... + class B: ... + """, + ), + ( # a comment after the literal is kept by the rewrite + """\ + __all__ = ["A"] # public API + + class A: ... + class B: ... + """, + """\ + __all__ = ["A", "B"] # public API + + class A: ... + class B: ... + """, + ), ]