Don't drop comments inside __all__, and only read markers from comments

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
 

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:

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

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)

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"]

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: ...
+        """,
+    ),
 ]