Better static analysis support

.pre-commit-config.yaml

@@ -80,6 +80,17 @@ repos:
           - "flake8-import-order<0.19,>=0.18"
           - flake8-print>=3.1.4,<4
 
+      - id: mypy
+        name: Type check Python (mypy)
+        entry: mypy -p pynocular
+        language: python
+        types: [file, python]
+        additional_dependencies:
+          - mypy==0.931
+          - types-dataclasses==0.6.1
+          - pydantic~=1.9.0
+        pass_filenames: false
+
       - id: circleci-config
         name: Pack the CircleCI config
         entry: .githooks/pre-commit-scripts/circleci.sh

README.md

@@ -450,6 +450,82 @@ with patch_database_model(Org, models=orgs), patch_database_model(
     assert org_get.tech_owner.username == "bberkley"
 ```
 
+### Type checking and mypy support
+
+Pynocular is fully type hinted and supports mypy in a limited way. Using mypy
+with the Pydantic mypy plugin will provide type checking for Pynocular database
+models' usual pydantic fields, but by default mypy will not know about the
+DatabaseModel specific methods.
+
+e.g.
+```python
+from pynocular.database_model import database_model
+from pydantic import BaseModel
+
+@database_model(...)
+class MyModel(BaseModel):
+  my_field: str = ...
+
+model = MyModel(...) # mypy understands this type
+print(model.my_field) # mypy will not error on this
+print(model.non_existent_field) # mypy will error here
+model_list = await MyModel.get_list(...) # mypy will (incorrectly) error here
+```
+
+The false positives are due to a limitation of python's type hint system that
+doesn't allow for dynamic class modification. The current workaround is to use
+the `BaseModel_database_model_hint` class provided by Pynocular, which is
+identical at runtime to the Pydantic one, but with extra type information. For
+convenience, the unmodified `BaseModel` class is also available to be imported
+from the same place.
+
+Warning: if you use `BaseModel_database_model_hint` to create a Pydantic model
+without the `database_model` decorator, mypy will incorrectly believe the model
+will have the pynocular provided functionality.
+
+e.g.
+
+```python
+from pynocular.database_model import database_model, BaseModel, BaseModel_database_model_hint
+
+@database_model(...)
+class MyModel_(BaseModel_database_model_hint):
+  my_field: str = ...
+
+model = MyModel(...) # mypy understands this type
+print(model.my_field) # mypy will not error on this
+print(model.non_existent_field) # mypy will error here
+model_list = await MyModel.get_list(...) # mypy will not error anymore
+print(model_list[0].non_existent_field) # mypy knows the type of model_list is
+                                        # List[MyModel], so catches the missing attribute
+model_list[0].my_field = "foo" # mypy doesn't error
+await model_list[0].save(some_bad_arg) # mypy errors on the incorrect call to save
+
+class MyNonDbModel(BaseModel):
+  my_other_field: str = ...
+
+other_list = await MyNonDbModel.create_list() # Mypy correctly errors on
+                                              # attempted use of pynocular provided method
+
+```
+
+In the future, this functionality may be rolled into a mypy plugin to avoid
+requiring the separate `BaseModel` imports.
+
+#### Nested Models
+
+Currently, Pynocular doesn't support nested models with mypy. As a workaround,
+you can append `# type: ignore` to lines using `nested_model` as a type
+annotation:
+
+```python
+
+class MyModel(BaseModel):
+  nested_db_field: nested_model(...) = ... # type: ignore
+
+```
+
+
 ## Development
 
 To develop Pynocular, install dependencies and enable the pre-commit hook.

pynocular/__init__.py

@@ -1,6 +1,6 @@
 """Lightweight ORM that lets you query your database using Pydantic models and asyncio"""
 
-__version__ = "0.18.0"
+__version__ = "0.19.0"
 
 from pynocular.database_model import DatabaseModel, UUID_STR
 from pynocular.engines import DatabaseType, DBInfo

pynocular/aiopg_transaction.py

@@ -1,7 +1,7 @@
 """Module for aiopg transaction utils"""
 import asyncio
 import sys
-from typing import Dict, Optional, Union
+from typing import Any, Dict, no_type_check, Optional, Type, Union
 
 import aiocontextvars as contextvars
 from aiopg.sa.connection import SAConnection
@@ -10,6 +10,7 @@
 transaction_connections_var = contextvars.ContextVar("transaction_connections")
 
 
+@no_type_check
 def get_current_task() -> asyncio.Task:
     """Get the current task when this method is called
 
@@ -37,12 +38,12 @@ def __init__(self, connection: SAConnection) -> None:
         self._conn = connection
         self.lock = asyncio.Lock()
 
-    async def execute(self, *args, **kwargs):
+    async def execute(self, *args: Any, **kwargs: Any) -> Any:
         """Wrapper around the `execute` method of the wrapped SAConnection"""
         async with self.lock:
             return await self._conn.execute(*args, **kwargs)
 
-    def __getattr__(self, attr):
+    def __getattr__(self, attr: str) -> Any:
         """Except for execute, all other attributes should pass through"""
         return getattr(self._conn, attr)
 
@@ -192,7 +193,10 @@ async def __aenter__(self) -> LockedConnection:
                 raise
         return conn
 
-    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+    @no_type_check
+    async def __aexit__(
+        self, exc_type: Optional[Type], exc_value: Any, tb: Any
+    ) -> None:
         """Exit the transaction context
 
         If this is the top level context, then commit the transaction (unless
@@ -251,7 +255,7 @@ def __init__(self, engine: aiopg.sa.engine.Engine) -> None:
         """
         super().__init__(engine)
         # The connection object, if functioning as standard connection
-        self._conn = None
+        self._conn: Optional[SAConnection] = None
 
     async def __aenter__(self) -> Union[LockedConnection, SAConnection]:
         """Conditionally establish the transaction context
@@ -269,7 +273,9 @@ async def __aenter__(self) -> Union[LockedConnection, SAConnection]:
         self._conn = await self._engine.acquire()
         return self._conn
 
-    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+    async def __aexit__(
+        self, exc_type: Optional[Type], exc_value: Any, tb: Any
+    ) -> None:
         """Exit the transaction context"""
         if self._conn is not None:
             await self._conn.close()