Add types to infobase db/server modules

Author: cdrini

infogami/infobase/_dbstore/store.py

@@ -39,12 +39,12 @@ def index(self, doc):
 class Store:
     """JSON Store."""
 
-    def __init__(self, db):
+    def __init__(self, db: web.DB):
         self.db = db
         self.indexer = StoreIndexer()
         self.listener = None
 
-    def get_row(self, key, for_update=False):
+    def get_row(self, key: str, for_update: bool = False):
         q = "SELECT * FROM store WHERE key=$key"
         if for_update:
             q += " FOR UPDATE NOWAIT"

infogami/infobase/client.py

@@ -69,10 +69,10 @@ class Connection:
     def __init__(self):
         self.auth_token = None
 
-    def set_auth_token(self, token):
+    def set_auth_token(self, token: str) -> None:
         self.auth_token = token
 
-    def get_auth_token(self):
+    def get_auth_token(self) -> str | None:
         return self.auth_token
 
     def request(self, sitename, path, method='GET', data=None):
@@ -223,11 +223,11 @@ def __iter__(self):
 
 
 class Site:
-    def __init__(self, conn, sitename):
+    def __init__(self, conn: Connection, sitename: str):
         self._conn = conn
         self.name = sitename
         # cache for storing pages requested in this HTTP request
-        self._cache = {}
+        self._cache: dict[tuple[str, int | None], web.storage] = {}
 
         self.store = Store(conn, sitename)
         self.seq = Sequence(conn, sitename)

infogami/infobase/core.py

@@ -165,15 +165,15 @@ class Store:
     Store manages one or many SiteStores.
     """
 
-    def create(self, sitename):
+    def create(self, sitename: str) -> "SiteStore":
         """Creates a new site with the given name and returns store for it."""
         raise NotImplementedError
 
-    def get(self, sitename):
+    def get(self, sitename: str) -> "SiteStore | None":
         """Returns store object for the given sitename."""
         raise NotImplementedError
 
-    def delete(self, sitename):
+    def delete(self, sitename: str):
         """Deletes the store for the specified sitename."""
         raise NotImplementedError
 

infogami/infobase/dbstore.py

@@ -3,6 +3,7 @@
 import datetime
 import logging
 import time
+from typing import TYPE_CHECKING, cast
 
 import web
 
@@ -12,7 +13,10 @@
 from infogami.infobase._dbstore.indexer import Indexer
 from infogami.infobase._dbstore.read import RecentChanges, get_bot_users
 from infogami.infobase._dbstore.save import PropertyManager, SaveImpl
-from infogami.infobase._dbstore.schema import Schema  # noqa: F401
+from infogami.infobase._dbstore.schema import Schema
+
+if TYPE_CHECKING:
+    from infogami.infobase.cache import Cache
 
 default_schema = None
 
@@ -25,21 +29,21 @@ def process_json(key, json_data):
 
 
 class DBSiteStore(common.SiteStore):
-    def __init__(self, db, schema):
+    def __init__(self, db: web.DB, schema: Schema):
         self.db = db
         self.schema = schema
         self.sitename = None
         self.indexer = Indexer()
         self.store = store.Store(self.db)
         self.seq = sequence.SequenceImpl(self.db)
 
-        self.cache = None
+        self.cache: "Cache | None" = None
         self.property_manager = PropertyManager(self.db)
 
     def get_store(self):
         return self.store
 
-    def set_cache(self, cache):
+    def set_cache(self, cache: "Cache"):
         self.cache = cache
 
     def get_metadata(self, key, for_update=False):
@@ -657,9 +661,9 @@ class DBStore(common.Store):
     It always returns a the same site irrespective of the sitename.
     """
 
-    def __init__(self, schema):
+    def __init__(self, schema: Schema):
         self.schema = schema
-        self.sitestore = None
+        self.sitestore: DBSiteStore | None = None
         self.db = create_database(**web.config.db_parameters)
 
     def has_initialized(self):
@@ -669,7 +673,7 @@ def has_initialized(self):
         except Exception:
             return False
 
-    def create(self, sitename):
+    def create(self, sitename: str) -> DBSiteStore:
         if self.sitestore is None:
             self.sitestore = DBSiteStore(self.db, self.schema)
             if not self.has_initialized():
@@ -678,7 +682,7 @@ def create(self, sitename):
         self.sitestore.initialize()
         return self.sitestore
 
-    def get(self, sitename):
+    def get(self, sitename: str) -> DBSiteStore | None:
         if self.sitestore is None:
             sitestore = DBSiteStore(self.db, self.schema)
             if not self.has_initialized():
@@ -699,9 +703,9 @@ def delete(self, sitename):
 class MultiDBStore(DBStore):
     """DBStore that works with multiple sites."""
 
-    def __init__(self, schema):
+    def __init__(self, schema: Schema):
         self.schema = schema
-        self.sitestores = {}
+        self.sitestores: dict[str, MultiDBSiteStore] = {}
         self.db = create_database(**web.config.db_parameters)
 
     def create(self, sitename):
@@ -783,8 +787,8 @@ def delete(self):
         pass
 
 
-def create_database(**params):
-    db = web.database(**params)
+def create_database(**params) -> web.DB:
+    db = cast("web.DB", web.database(**params))
 
     # monkey-patch query method to collect stats
     _query = db.query

infogami/infobase/infobase.py

@@ -7,6 +7,8 @@
 
 import datetime
 import json
+from collections.abc import Callable
+from typing import TYPE_CHECKING
 
 import web
 
@@ -20,48 +22,51 @@
     writequery,
 )
 
+if TYPE_CHECKING:
+    from infogami.infobase.dbstore import DBSiteStore, DBStore
+
 
 class Infobase:
     """Infobase contains multiple sites."""
 
-    def __init__(self, store, secret_key):
+    def __init__(self, store: "DBStore", secret_key: str):
         self.store = store
         self.secret_key = secret_key
-        self.sites = {}
-        self.event_listeners = []
+        self.sites: "dict[str, Site]" = {}
+        self.event_listeners: "list[Callable]" = []
 
         if config.startup_hook:
             config.startup_hook(self)
 
-    def create(self, sitename):
+    def create(self, sitename: str) -> "Site":
         """Creates a new site with the sitename."""
         site = Site(self, sitename, self.store.create(sitename), self.secret_key)
         site.bootstrap()
         self.sites[sitename] = site
         return site
 
-    def get(self, sitename):
+    def get(self, sitename: str) -> "Site | None":
         """Returns the site with the given name."""
         if sitename in self.sites:
             site = self.sites[sitename]
         else:
             store = self.store.get(sitename)
             if store is None:
                 return None
-            site = Site(self, sitename, self.store.get(sitename), self.secret_key)
+            site = Site(self, sitename, store, self.secret_key)
             self.sites[sitename] = site
         return site
 
-    def delete(self, sitename):
+    def delete(self, sitename: str):
         """Deletes the site with the given name."""
         if sitename in self.sites:
             del self.sites[sitename]
         return self.store.delete(sitename)
 
-    def add_event_listener(self, listener):
+    def add_event_listener(self, listener: Callable):
         self.event_listeners.append(listener)
 
-    def remove_event_listener(self, listener):
+    def remove_event_listener(self, listener: Callable):
         try:
             self.event_listeners.remove(listener)
         except ValueError:
@@ -78,15 +83,21 @@ def fire_event(self, event):
 class Site:
     """A site of infobase."""
 
-    def __init__(self, _infobase, sitename, store, secret_key):
+    def __init__(
+        self,
+        _infobase: Infobase,
+        sitename: str,
+        store: "DBSiteStore",
+        secret_key: str,
+    ):
         self._infobase = _infobase
         self.sitename = sitename
         self.store = store
         self.cache = cache.Cache()
         self.store.set_cache(self.cache)
         self.account_manager = account.AccountManager(self, secret_key)
 
-        self._triggers = {}
+        self._triggers: dict[str, list[Callable]] = {}
         store.store.set_listener(self._log_store_action)
         store.seq.set_listener(self._log_store_action)
 

infogami/infobase/server.py

@@ -187,10 +187,10 @@ def from_json(s):
         raise common.BadData(message="Bad JSON: " + str(e))
 
 
-_infobase = None
+_infobase: infobase.Infobase | None = None
 
 
-def get_site(sitename):
+def get_site(sitename: str) -> infobase.Site | None:
     global _infobase
     if not _infobase:
         schema = dbstore.default_schema or dbstore.Schema()