refactor: fallback to pyi file generation

Signed-off-by: Gabor Boros <[email protected]>

Author: gabor-boros

Makefile

@@ -94,6 +94,12 @@ docs: ## generate Sphinx HTML documentation, including API docs
 	poetry run $(MAKE) -C docs html
 	$(BROWSER) docs/_build/html/index.html
 
+.PHONY: generate-init-pyi
+generate-init-pyi: ## generate the __init__.pyi file
+	poetry run python scripts/generate_init_pyi.py && \
+	poetry run isort rethinkdb/__init__.pyi && \
+	poetry run black rethinkdb/__init__.pyi
+
 .PHONY: format
 format: ## run formatters on the package
 	poetry run isort rethinkdb tests

rethinkdb/__init__.py

@@ -15,224 +15,10 @@
 # This file incorporates work covered by the following copyright:
 # Copyright 2010-2016 RethinkDB, all rights reserved.
 
-# pylint: disable=redefined-builtin, unused-import
 
-from types import SimpleNamespace
-import warnings
-
-from rethinkdb import net
-from rethinkdb.ast import expr
-from rethinkdb.query import (
-    add,
-    and_,
-    april,
-    args,
-    asc,
-    august,
-    avg,
-    binary,
-    bit_and,
-    bit_not,
-    bit_or,
-    bit_sal,
-    bit_sar,
-    bit_xor,
-    branch,
-    ceil,
-    circle,
-    contains,
-    count,
-    db,
-    db_create,
-    db_drop,
-    db_list,
-    december,
-    desc,
-    distance,
-    distinct,
-    div,
-    do,
-    epoch_time,
-    eq,
-    error,
-    february,
-    floor,
-    format,
-    friday,
-    ge,
-    geojson,
-    grant,
-    group,
-    gt,
-    http,
-    info,
-    intersects,
-    iso8601,
-    january,
-    js,
-    json,
-    july,
-    june,
-    le,
-    line,
-    literal,
-    lt,
-    make_timezone,
-    map,
-    march,
-    max,
-    maxval,
-    may,
-    min,
-    minval,
-    mod,
-    monday,
-    mul,
-    ne,
-    not_,
-    november,
-    now,
-    object,
-    october,
-    or_,
-    point,
-    polygon,
-    random,
-    range,
-    reduce,
-    round,
-    row,
-    saturday,
-    september,
-    sub,
-    sum,
-    sunday,
-    table,
-    table_create,
-    table_drop,
-    table_list,
-    thursday,
-    time,
-    tuesday,
-    type_of,
-    union,
-    uuid,
-    wednesday,
-)
-
-# pylint: enable=redefined-builtin, unused-import
-
-__version__ = "2.5.0"
-
-# Create the r namespace object containing all query functions
-r = SimpleNamespace()
-
-query_functions = {
-    "add": add,
-    "and_": and_,
-    "april": april,
-    "args": args,
-    "asc": asc,
-    "august": august,
-    "avg": avg,
-    "binary": binary,
-    "bit_and": bit_and,
-    "bit_not": bit_not,
-    "bit_or": bit_or,
-    "bit_sal": bit_sal,
-    "bit_sar": bit_sar,
-    "bit_xor": bit_xor,
-    "branch": branch,
-    "ceil": ceil,
-    "circle": circle,
-    "contains": contains,
-    "count": count,
-    "db": db,
-    "db_create": db_create,
-    "db_drop": db_drop,
-    "db_list": db_list,
-    "december": december,
-    "desc": desc,
-    "distance": distance,
-    "distinct": distinct,
-    "div": div,
-    "do": do,
-    "epoch_time": epoch_time,
-    "eq": eq,
-    "expr": expr,
-    "error": error,
-    "february": february,
-    "floor": floor,
-    "format": format,
-    "friday": friday,
-    "ge": ge,
-    "geojson": geojson,
-    "grant": grant,
-    "group": group,
-    "gt": gt,
-    "http": http,
-    "info": info,
-    "intersects": intersects,
-    "iso8601": iso8601,
-    "january": january,
-    "json": json,
-    "july": july,
-    "june": june,
-    "le": le,
-    "line": line,
-    "literal": literal,
-    "lt": lt,
-    "make_timezone": make_timezone,
-    "map": map,
-    "march": march,
-    "max": max,
-    "maxval": maxval,
-    "may": may,
-    "min": min,
-    "minval": minval,
-    "mod": mod,
-    "monday": monday,
-    "mul": mul,
-    "ne": ne,
-    "not_": not_,
-    "november": november,
-    "now": now,
-    "object": object,
-    "october": october,
-    "or_": or_,
-    "point": point,
-    "polygon": polygon,
-    "random": random,
-    "range": range,
-    "reduce": reduce,
-    "round": round,
-    "row": row,
-    "saturday": saturday,
-    "september": september,
-    "sub": sub,
-    "sum": sum,
-    "sunday": sunday,
-    "table": table,
-    "table_create": table_create,
-    "table_drop": table_drop,
-    "table_list": table_list,
-    "thursday": thursday,
-    "time": time,
-    "tuesday": tuesday,
-    "type_of": type_of,
-    "union": union,
-    "uuid": uuid,
-    "wednesday": wednesday,
-    "js": js,
-}
-
-for name, func in query_functions.items():
-    setattr(r, name, func)
-
-
-class Client:
+class RethinkDB:
     """
-    Client is a wrapper around RethinkDB connection handling.
+    RethinkDB is a wrapper around RethinkDB connection handling.
 
     It constructs the connection handlers and event loops, re-exports internal modules for easier
     use, and sets the event loop.
@@ -241,9 +27,22 @@ class Client:
     def __init__(self):
         super().__init__()
 
+        # pylint: disable=import-outside-toplevel
+        from rethinkdb import ast, errors, net, query
+
+        self.ast = ast
+        self.errors = errors
+        self.net = net
+        self.query = query
+
         net.Connection._r = self
         self.connection_type = None
 
+        # Dynamically assign every re-exported internal module's function to self
+        for module in (self.net, self.query, self.ast, self.errors):
+            for function_name in module.__all__:
+                setattr(self, function_name, getattr(module, function_name))
+
         self.make_connection = net.make_connection
         self.set_loop_type(None)
 
@@ -253,8 +52,10 @@ def set_loop_type(self, library=None) -> None:
         """
 
         if library == "asyncio":
-            warnings.warn(f"{library} is not yet supported, using the default one")
-            library = None
+            # pylint: disable=import-outside-toplevel
+            from rethinkdb.net_asyncio import Connection as AsyncioConnection
+
+            self.connection_type = AsyncioConnection
 
         if library == "gevent":
             # pylint: disable=import-outside-toplevel
@@ -269,8 +70,10 @@ def set_loop_type(self, library=None) -> None:
             self.connection_type = TornadoConnection
 
         if library == "trio":
-            warnings.warn(f"{library} is not yet supported, using the default one")
-            library = None
+            # pylint: disable=import-outside-toplevel
+            from rethinkdb.net_trio import Connection as TrioConnection
+
+            self.connection_type = TrioConnection
 
         if library == "twisted":
             # pylint: disable=import-outside-toplevel
@@ -279,11 +82,17 @@ def set_loop_type(self, library=None) -> None:
             self.connection_type = TwistedConnection
 
         if library is None or self.connection_type is None:
-            self.connection_type = net.DefaultConnection
+            # pylint: disable=import-outside-toplevel
+            from rethinkdb.net import DefaultConnection
+
+            self.connection_type = DefaultConnection
 
     def connect(self, *connect_args, **kwargs):
         """
         Make a connection to the database.
         """
 
         return self.make_connection(self.connection_type, *connect_args, **kwargs)
+
+
+r = RethinkDB()