some typing (#286)

peewee_async/aio_model.py

@@ -87,7 +87,7 @@ class AioModelRaw(peewee.ModelRaw, AioQueryMixin):
 class AioSelectMixin(AioQueryMixin):
 
     @peewee.database_required
-    async def aio_scalar(self, database, as_tuple: bool = False):
+    async def aio_scalar(self, database: AioDatabase, as_tuple: bool = False) -> Any:
         """
         Get single value from ``select()`` query, i.e. for aggregation.
 
@@ -103,7 +103,7 @@ async def fetch_results(cursor):
 
         return rows[0] if rows and not as_tuple else rows
 
-    async def aio_get(self, database=None):
+    async def aio_get(self, database: Optional[AioDatabase] = None):
         """
         Async version of **peewee.SelectBase.get**
 
@@ -120,7 +120,7 @@ async def aio_get(self, database=None):
                                           (clone.model, sql, params))
 
     @peewee.database_required
-    async def aio_count(self, database, clear_limit: bool =False) -> int:
+    async def aio_count(self, database: AioDatabase, clear_limit: bool = False) -> int:
         """
         Async version of **peewee.SelectBase.count**
 
@@ -143,7 +143,7 @@ async def aio_count(self, database, clear_limit: bool =False) -> int:
         )
 
     @peewee.database_required
-    async def aio_exists(self, database):
+    async def aio_exists(self, database: AioDatabase) -> bool:
         """
         Async version of **peewee.SelectBase.exists**
 

peewee_async/databases.py

@@ -17,14 +17,14 @@ class AioDatabase(peewee.Database):
     pool_backend_cls: Type[PoolBackend]
     pool_backend: PoolBackend
 
-    def __init__(self, *args, **kwargs):
-        self.pool_params = {}
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self.pool_params: Dict[str, Any] = {}
         super().__init__(*args, **kwargs)
 
-    def init_pool_params_defaults(self):
+    def init_pool_params_defaults(self) -> None:
         pass
 
-    def init_pool_params(self):
+    def init_pool_params(self) -> None:
         self.init_pool_params_defaults()
         self.pool_params.update(
             {

tests/test_database.py

@@ -76,7 +76,7 @@ async def test_deferred_init(db_name: str) -> None:
 
 
 @pytest.mark.parametrize('db_name', PG_DBS + MYSQL_DBS)
-async def test_connections_param(db_name):
+async def test_connections_param(db_name: str) -> None:
     default_params = DB_DEFAULTS[db_name].copy()
     default_params['min_connections'] = 2
     default_params['max_connections'] = 3
@@ -85,38 +85,38 @@ async def test_connections_param(db_name):
     database = db_cls(**default_params)
     await database.aio_connect()
 
-    assert database.pool_backend.pool._minsize == 2
-    assert database.pool_backend.pool._free.maxlen == 3
+    assert database.pool_backend.pool._minsize == 2  # type: ignore
+    assert database.pool_backend.pool._free.maxlen == 3  # type: ignore
 
     await database.aio_close()
 
 
 @dbs_mysql
-async def test_mysql_params(db):
+async def test_mysql_params(db: AioDatabase) -> None:
     async with db.aio_connection() as connection_1:
-        assert connection_1.autocommit_mode is True
-    assert db.pool_backend.pool._recycle == 2
+        assert connection_1.autocommit_mode is True  # type: ignore
+    assert db.pool_backend.pool._recycle == 2  # type: ignore
 
 
 @pytest.mark.parametrize(
     "db",
     ["postgres-pool"], indirect=["db"]
 )
-async def test_pg_json_hstore__params(db):
+async def test_pg_json_hstore__params(db: AioDatabase) -> None:
     await db.aio_connect()
-    assert db.pool_backend.pool._enable_json is False
-    assert db.pool_backend.pool._enable_hstore is False
-    assert db.pool_backend.pool._timeout == 30
-    assert db.pool_backend.pool._recycle == 1.5
+    assert db.pool_backend.pool._enable_json is False  # type: ignore
+    assert db.pool_backend.pool._enable_hstore is False  # type: ignore
+    assert db.pool_backend.pool._timeout == 30  # type: ignore
+    assert db.pool_backend.pool._recycle == 1.5  # type: ignore
 
 
 @pytest.mark.parametrize(
     "db",
     ["postgres-pool-ext"], indirect=["db"]
 )
-async def test_pg_ext_json_hstore__params(db):
+async def test_pg_ext_json_hstore__params(db: AioDatabase) -> None:
     await db.aio_connect()
-    assert db.pool_backend.pool._enable_json is True
-    assert db.pool_backend.pool._enable_hstore is False
-    assert db.pool_backend.pool._timeout == 30
-    assert db.pool_backend.pool._recycle == 1.5
+    assert db.pool_backend.pool._enable_json is True  # type: ignore
+    assert db.pool_backend.pool._enable_hstore is False  # type: ignore
+    assert db.pool_backend.pool._timeout == 30  # type: ignore
+    assert db.pool_backend.pool._recycle == 1.5  # type: ignore