fix docs generation

Author: DenisBiryukov91

docs/stubs_to_sources.py

@@ -28,6 +28,7 @@
 
 PACKAGE = (Path(__file__) / "../../zenoh").resolve()
 __INIT__ = PACKAGE / "__init__.py"
+EXT = PACKAGE / "ext.py"
 
 
 def _unstable(item):
@@ -74,7 +75,8 @@ def visit_FunctionDef(self, node: ast.FunctionDef):
                 node.decorator_list.clear()
                 if node.name not in ("recv", "try_recv", "__iter__"):
                     # retrieve the handled type (Scout/Reply/etc.) from the return type
-                    assert isinstance(node.returns, ast.Subscript)
+                    if not isinstance(node.returns, ast.Subscript):
+                        continue
                     if isinstance(node.returns.slice, ast.Subscript):
                         # `Subscriber[Handler[Sample]]` case
                         tp = node.returns.slice.slice
@@ -98,23 +100,26 @@ def visit_FunctionDef(self, node: ast.FunctionDef):
 
 
 def main():
-    # remove __init__.pyi
-    __INIT__.unlink()
+    fnames = [__INIT__, EXT]
+    for fname in fnames:
+        # remove *.py
+        fname.unlink()
     # rename stubs
     for entry in PACKAGE.glob("*.pyi"):
         entry.rename(PACKAGE / f"{entry.stem}.py")
-    # read stub code
-    with open(__INIT__) as f:
-        stub: ast.Module = ast.parse(f.read())
-        # replace _unstable
-        for i, stmt in enumerate(stub.body):
-            if isinstance(stmt, ast.FunctionDef) and stmt.name == "_unstable":
-                stub.body[i] = ast.parse(inspect.getsource(_unstable))
-        # remove overload
-        stub = RemoveOverload().visit(stub)
-    # write modified code
-    with open(__INIT__, "w") as f:
-        f.write(ast.unparse(stub))
+    for fname in fnames:
+        # read stub code
+        with open(fname) as f:
+            stub: ast.Module = ast.parse(f.read())
+            # replace _unstable
+            for i, stmt in enumerate(stub.body):
+                if isinstance(stmt, ast.FunctionDef) and stmt.name == "_unstable":
+                    stub.body[i] = ast.parse(inspect.getsource(_unstable))
+            # remove overload
+            stub = RemoveOverload().visit(stub)
+        # write modified code
+        with open(fname, "w") as f:
+            f.write(ast.unparse(stub))
 
 
 if __name__ == "__main__":