From cea9fc6c7749d3a13d66af1c4e221292640a9d8e Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Thu, 2 Nov 2023 12:43:06 +0100
Subject: [PATCH] Ensure __all__ are consistent

---
 .../telethon/_impl/client/client/__init__.py  |  6 +-
 .../telethon/_impl/client/parsers/__init__.py |  4 +-
 .../_impl/client/types/buttons/__init__.py    |  1 -
 .../_impl/client/types/chat/__init__.py       | 11 +---
 client/src/telethon/_impl/session/__init__.py |  2 +-
 client/src/telethon/_impl/tl/__init__.py      | 11 +---
 tools/copy_init_imports.py                    | 61 +++++++++++++++++++
 7 files changed, 67 insertions(+), 29 deletions(-)
 create mode 100644 tools/copy_init_imports.py

diff --git a/client/src/telethon/_impl/client/client/__init__.py b/client/src/telethon/_impl/client/client/__init__.py
index 21b01780..2094dad2 100644
--- a/client/src/telethon/_impl/client/client/__init__.py
+++ b/client/src/telethon/_impl/client/client/__init__.py
@@ -1,8 +1,4 @@
 from .client import Client
 from .net import Config
 
-__all__ = [
-    "InlineResult",
-    "Client",
-    "Config",
-]
+__all__ = ["Client", "Config"]
diff --git a/client/src/telethon/_impl/client/parsers/__init__.py b/client/src/telethon/_impl/client/parsers/__init__.py
index 4d3985fb..3a89098e 100644
--- a/client/src/telethon/_impl/client/parsers/__init__.py
+++ b/client/src/telethon/_impl/client/parsers/__init__.py
@@ -4,8 +4,8 @@ from .markdown import parse as parse_markdown_message
 from .markdown import unparse as generate_markdown_message
 
 __all__ = [
-    "generate_html_message",
     "parse_html_message",
-    "generate_markdown_message",
+    "generate_html_message",
     "parse_markdown_message",
+    "generate_markdown_message",
 ]
diff --git a/client/src/telethon/_impl/client/types/buttons/__init__.py b/client/src/telethon/_impl/client/types/buttons/__init__.py
index 79918db0..221734d8 100644
--- a/client/src/telethon/_impl/client/types/buttons/__init__.py
+++ b/client/src/telethon/_impl/client/types/buttons/__init__.py
@@ -85,5 +85,4 @@ __all__ = [
     "SwitchInline",
     "Text",
     "Url",
-    "create",
 ]
diff --git a/client/src/telethon/_impl/client/types/chat/__init__.py b/client/src/telethon/_impl/client/types/chat/__init__.py
index 1d03aea7..21c860a8 100644
--- a/client/src/telethon/_impl/client/types/chat/__init__.py
+++ b/client/src/telethon/_impl/client/types/chat/__init__.py
@@ -80,13 +80,4 @@ def expand_peer(peer: abcs.Peer, *, broadcast: Optional[bool]) -> Chat:
         raise RuntimeError("unexpected case")
 
 
-__all__ = [
-    "Chat",
-    "ChatLike",
-    "Channel",
-    "Group",
-    "User",
-    "build_chat_map",
-    "peer_id",
-    "expand_peer",
-]
+__all__ = ["Channel", "Chat", "Group", "User"]
diff --git a/client/src/telethon/_impl/session/__init__.py b/client/src/telethon/_impl/session/__init__.py
index e0d4ff31..f553a210 100644
--- a/client/src/telethon/_impl/session/__init__.py
+++ b/client/src/telethon/_impl/session/__init__.py
@@ -21,6 +21,7 @@ __all__ = [
     "NO_UPDATES_TIMEOUT",
     "USER_CHANNEL_DIFF_LIMIT",
     "Gap",
+    "MessageBox",
     "PossibleGap",
     "PrematureEndReason",
     "PtsInfo",
@@ -30,7 +31,6 @@ __all__ = [
     "Session",
     "UpdateState",
     "User",
-    "MessageBox",
     "MemorySession",
     "SqliteSession",
     "Storage",
diff --git a/client/src/telethon/_impl/tl/__init__.py b/client/src/telethon/_impl/tl/__init__.py
index 89c247e4..bc66da10 100644
--- a/client/src/telethon/_impl/tl/__init__.py
+++ b/client/src/telethon/_impl/tl/__init__.py
@@ -2,13 +2,4 @@ from . import abcs, functions, mtproto, types
 from .core import Request
 from .layer import LAYER, TYPE_MAPPING
 
-__all__ = [
-    "abcs",
-    "core",
-    "functions",
-    "mtproto",
-    "types",
-    "Request",
-    "LAYER",
-    "TYPE_MAPPING",
-]
+__all__ = ['abcs', 'functions', 'mtproto', 'types', 'Request', 'LAYER', 'TYPE_MAPPING']
diff --git a/tools/copy_init_imports.py b/tools/copy_init_imports.py
new file mode 100644
index 00000000..43da2ea3
--- /dev/null
+++ b/tools/copy_init_imports.py
@@ -0,0 +1,61 @@
+"""
+Scan the `client/` directory for __init__.py files.
+For every depth-1 import, add it, in order, to the __all__ variable.
+"""
+import ast
+import os
+import re
+from pathlib import Path
+from typing import List
+
+
+class ImportVisitor(ast.NodeVisitor):
+    def __init__(self) -> None:
+        self.imported_names: List[str] = []
+
+    def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        if node.level == 1:
+            for name in node.names:
+                self.imported_names.append(name.asname or name.name)
+
+
+def main() -> None:
+    impl_root = Path.cwd() / "client/src/telethon/_impl"
+    autogenerated_re = re.compile(
+        rf"(tl|mtproto){re.escape(os.path.sep)}(abcs|functions|types)"
+    )
+
+    files = []
+    for file in impl_root.rglob("__init__.py"):
+        file_str = str(file)
+        if autogenerated_re.search(file_str):
+            continue
+
+        files.append(file_str)
+
+        with file.open(encoding="utf-8") as fd:
+            contents = fd.read()
+        lines = contents.splitlines(True)
+
+        module = ast.parse(contents)
+
+        visitor = ImportVisitor()
+        visitor.visit(module)
+
+        for stmt in module.body:
+            match stmt:
+                case ast.Assign(targets=[ast.Name(id="__all__")], value=ast.List()):
+                    # Not rewriting the AST to preserve comments and formatting.
+                    lines[stmt.lineno - 1 : stmt.end_lineno] = [
+                        "__all__ = ["
+                        + ", ".join(map(repr, visitor.imported_names))
+                        + "]\n"
+                    ]
+                    break
+
+        with file.open("w", encoding="utf-8", newline="\n") as fd:
+            fd.writelines(lines)
+
+
+if __name__ == "__main__":
+    main()