From 1acbeb217bca6d7f40edce3094259101ff29690d Mon Sep 17 00:00:00 2001
From: Vadim Levin <vadim.levin@xperience.ai>
Date: Thu, 15 Jun 2023 16:32:48 +0300
Subject: [PATCH] feat: re-export symbols to cv2 level

- Re-export native submodules of cv2 package level.
- Re-export  manually registered  symbols like cv2.mat_wrapper.Mat
---
 .../typing_stubs_generation/generation.py     | 52 ++++++++++++++++++-
 .../nodes/namespace_node.py                   | 22 +++++---
 2 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/modules/python/src2/typing_stubs_generation/generation.py b/modules/python/src2/typing_stubs_generation/generation.py
index 018c55414a..5c32f43ff9 100644
--- a/modules/python/src2/typing_stubs_generation/generation.py
+++ b/modules/python/src2/typing_stubs_generation/generation.py
@@ -70,10 +70,11 @@ def generate_typing_stubs(root: NamespaceNode, output_path: Path):
     # checked and at least 1 node is still unresolved.
     root.resolve_type_nodes()
     _generate_typing_module(root, output_path)
+    _populate_reexported_symbols(root)
     _generate_typing_stubs(root, output_path)
 
 
-def _generate_typing_stubs(root: NamespaceNode, output_path: Path):
+def _generate_typing_stubs(root: NamespaceNode, output_path: Path) -> None:
     output_path = Path(output_path) / root.export_name
     output_path.mkdir(parents=True, exist_ok=True)
 
@@ -85,6 +86,8 @@ def _generate_typing_stubs(root: NamespaceNode, output_path: Path):
     # Write required imports at the top of file
     _write_required_imports(required_imports, output_stream)
 
+    _write_reexported_symbols_section(root, output_stream)
+
     # Write constants section, because constants don't impose any dependencies
     _generate_section_stub(StubSection("# Constants", ConstantNode), root,
                            output_stream, 0)
@@ -582,6 +585,53 @@ def _collect_required_imports(root: NamespaceNode) -> Set[str]:
     return required_imports
 
 
+def _populate_reexported_symbols(root: NamespaceNode) -> None:
+    # Re-export all submodules to allow referencing symbols in submodules
+    # without submodule import. Example:
+    # `cv2.aruco.ArucoDetector` should be accessible without `import cv2.aruco`
+    for submodule in root.namespaces.values():
+        root.reexported_submodules.append(submodule.export_name)
+
+    # Special cases, symbols defined in possible pure Python submodules should be
+    root.reexported_submodules_symbols["mat_wrapper"].append("Mat")
+
+def _write_reexported_symbols_section(module: NamespaceNode, output_stream: StringIO) -> None:
+    """Write re-export section for the given module.
+
+    Re-export statements have from `from module_name import smth as smth`.
+    Example:
+    ```python
+    from cv2 import aruco as aruco
+    from cv2 import cuda as cuda
+    from cv2 import ml as ml
+    from cv2.mat_wrapper import Mat as Mat
+    ```
+
+    Args:
+        module (NamespaceNode): Module with re-exported symbols.
+        output_stream (StringIO): Output stream for re-export statements.
+    """
+
+    parent_name = module.full_export_name
+    for submodule in sorted(module.reexported_submodules):
+        output_stream.write(
+            "from {0} import {1} as {1}\n".format(parent_name, submodule)
+        )
+
+    for submodule, symbols in sorted(module.reexported_submodules_symbols.items(),
+                                     key=lambda kv: kv[0]):
+        for symbol in symbols:
+            output_stream.write(
+                "from {0}.{1} import {2} as {2}\n".format(
+                    parent_name, submodule, symbol
+                )
+            )
+
+    if len(module.reexported_submodules) or \
+            len(module.reexported_submodules_symbols):
+        output_stream.write("\n\n")
+
+
 def _write_required_imports(required_imports: Collection[str],
                             output_stream: StringIO) -> None:
     """Writes all entries of `required_imports` to the `output_stream`.
diff --git a/modules/python/src2/typing_stubs_generation/nodes/namespace_node.py b/modules/python/src2/typing_stubs_generation/nodes/namespace_node.py
index 2206f0e586..8d0d04b5a5 100644
--- a/modules/python/src2/typing_stubs_generation/nodes/namespace_node.py
+++ b/modules/python/src2/typing_stubs_generation/nodes/namespace_node.py
@@ -1,14 +1,13 @@
-from typing import Type, Iterable, Sequence, Tuple, Optional, Dict
 import itertools
 import weakref
-
-from .node import ASTNode, ASTNodeType
+from collections import defaultdict
+from typing import Dict, List, Optional, Sequence, Tuple, Type
 
 from .class_node import ClassNode, ClassProperty
-from .function_node import FunctionNode
-from .enumeration_node import EnumerationNode
 from .constant_node import ConstantNode
-
+from .enumeration_node import EnumerationNode
+from .function_node import FunctionNode
+from .node import ASTNode, ASTNodeType
 from .type_node import TypeResolutionError
 
 
@@ -18,6 +17,17 @@ class NamespaceNode(ASTNode):
     NamespaceNode can have other namespaces, classes, functions, enumerations
     and global constants as its children nodes.
     """
+    def __init__(self, name: str, parent: Optional[ASTNode] = None,
+                 export_name: Optional[str] = None) -> None:
+        super().__init__(name, parent, export_name)
+        self.reexported_submodules: List[str] = []
+        """List of reexported submodules"""
+
+        self.reexported_submodules_symbols: Dict[str, List[str]] = defaultdict(list)
+        """Mapping between submodules export names and their symbols re-exported
+        in this module"""
+
+
     @property
     def node_type(self) -> ASTNodeType:
         return ASTNodeType.Namespace