(fix): Change raw/base docs types generation to get proper typehints

(KurimuzonAkuma/kurigram#151)

Follow-Up: bf83f57

Co-authored-by: Danipulok <[email protected]>
Co-authored-by: KurimuzonAkuma <[email protected]>

Author: 3 people

compiler/docs/compiler.py

@@ -21,6 +21,9 @@
 import re
 import shutil
 
+from dataclasses import dataclass
+from typing import Literal, Optional
+
 HOME = "compiler/docs"
 DESTINATION = "docs/source/telegram"
 PYROGRAM_API_DEST = "docs/source/api"
@@ -39,6 +42,64 @@ def snek(s: str):
     return re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", s).lower()
 
 
+def _extract_union_name(node: ast.AST) -> Optional[str]:
+    """Extract the name of a variable that is assigned a Union type.
+
+    :param node: The AST node to extract the variable name from.
+    :return: The variable name if it is assigned a Union type, otherwise None.
+
+    >>> import ast
+    >>> parsed_ast = ast.parse("User = Union[raw.types.UserEmpty]")
+    >>> _extract_union_name(parsed_ast.body[0])
+    'User'
+    """
+
+    # Check if the assigned value is a Union type
+    if isinstance(node, ast.Assign) and isinstance(node.value, ast.Subscript):
+        if isinstance(node.value.value, ast.Name) and node.value.value.id == "Union":
+            # Extract variable name
+            if isinstance(node.targets[0], ast.Name):
+                return node.targets[0].id  # Variable name
+
+
+def _extract_class_name(node: ast.AST) -> Optional[str]:
+    """Extract the name of a class.
+
+    :param node: The AST node to extract the class name from.
+    :return: The class name if it is a class, otherwise None.
+
+    >>> import ast
+    >>> parsed_ast = ast.parse("class User: pass")
+    >>> _extract_class_name(parsed_ast.body[0])
+    'User'
+    """
+
+    if isinstance(node, ast.ClassDef):
+        return node.name  # Class name
+
+
+NodeType = Literal["class", "union"]
+
+
+@dataclass
+class NodeInfo:
+    name: str
+    type: NodeType
+
+
+def parse_node_info(node: ast.AST) -> Optional[NodeInfo]:
+    """Parse an AST node and extract the class or variable name."""
+    class_name = _extract_class_name(node)
+    if class_name:
+        return NodeInfo(name=class_name, type="class")
+
+    union_name = _extract_union_name(node)
+    if union_name:
+        return NodeInfo(name=union_name, type="union")
+
+    return None
+
+
 def generate(source_path, base):
     all_entities = {}
 
@@ -54,13 +115,13 @@ def build(path, level=0):
                     p = ast.parse(f.read())
 
                 for node in ast.walk(p):
-                    if isinstance(node, ast.ClassDef):
-                        name = node.name
+                    node_info = parse_node_info(node)
+                    if node_info:
                         break
                 else:
                     continue
 
-                full_path = os.path.basename(path) + "/" + snek(name).replace("_", "-") + ".rst"
+                full_path = os.path.basename(path) + "/" + snek(node_info.name).replace("_", "-") + ".rst"
 
                 if level:
                     full_path = base + "/" + full_path
@@ -69,25 +130,41 @@ def build(path, level=0):
                 if namespace in ["base", "types", "functions"]:
                     namespace = ""
 
-                full_name = f"{(namespace + '.') if namespace else ''}{name}"
+                full_name = f"{(namespace + '.') if namespace else ''}{node_info.name}"
 
                 os.makedirs(os.path.dirname(DESTINATION + "/" + full_path), exist_ok=True)
 
                 with open(DESTINATION + "/" + full_path, "w", encoding="utf-8") as f:
+                    title_markup = "=" * len(full_name)
+                    full_class_path = "pyrogram.raw.{}".format(
+                        ".".join(full_path.split("/")[:-1]) + "." + node_info.name
+                    )
+                    if node_info.type == "class":
+                        directive_type = "autoclass"
+                        directive_suffix = "()"
+                        directive_option = "members"
+                    elif node_info.type == "union":
+                        directive_type = "autodata"
+                        directive_suffix = ""
+                        directive_option = "annotation"
+                    else:
+                        raise ValueError(f"Unknown node type: `{node_info.type}`")
+
                     f.write(
                         page_template.format(
                             title=full_name,
-                            title_markup="=" * len(full_name),
-                            full_class_path="pyrogram.raw.{}".format(
-                                ".".join(full_path.split("/")[:-1]) + "." + name
-                            )
+                            title_markup=title_markup,
+                            directive_type=directive_type,
+                            full_class_path=full_class_path,
+                            directive_suffix=directive_suffix,
+                            directive_option=directive_option,
                         )
                     )
 
                 if last not in all_entities:
                     all_entities[last] = []
 
-                all_entities[last].append(name)
+                all_entities[last].append(node_info.name)
 
     build(source_path)