More typing

Author: joente

.github/workflows/ci.yml

@@ -13,19 +13,22 @@ jobs:
     strategy:
       max-parallel: 4
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install pytest pycodestyle
+          pip install pytest pycodestyle pyright
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
       - name: Lint with PyCodeStyle
         run: |
           find . -name \*.py -exec pycodestyle {} +
+      - name: PyRight
+        run: |
+          pyright

aiogcd/connector/__init__.py

@@ -3,6 +3,6 @@
 Created on: May 19, 2017
     Author: Jeroen van der Heijden <jeroen@cesbit.com>
 """
-from .connector import GcdConnector, GcdServiceAccountConnector
-from .client_token import Token
-from .service_account_token import ServiceAccountToken
+from .connector import GcdConnector, GcdServiceAccountConnector  # noqa: F401
+from .client_token import Token  # noqa: F401
+from .service_account_token import ServiceAccountToken  # noqa: F401

aiogcd/connector/client_token.py

@@ -95,6 +95,7 @@ async def get(self) -> str:
 
         :return: Access token (string)
         """
+        assert self._refresh_ts is not None and self._token is not None
         async with self._lock:
             if self._refresh_ts < time.time():
                 await self._refresh_token()
@@ -127,6 +128,7 @@ async def connect(self):
                 logging.info('Token is valid.')
 
     async def _refresh_token(self):
+        assert self._token is not None
         logging.info(
             'Token has exceeded half of the expiration time, '
             'starting a token refresh.')

aiogcd/connector/connector.py

@@ -7,7 +7,7 @@
 import os
 import json
 import aiohttp
-from typing import Iterable, Optional, Any, Union
+from typing import Iterable, Any
 from .client_token import Token
 from .service_account_token import ServiceAccountToken
 from .entity import Entity
@@ -43,7 +43,7 @@ def __init__(
             client_secret: str,
             token_file: str,
             scopes: Iterable[str] = DEFAULT_SCOPES,
-            namespace_id: Optional[str] = None):
+            namespace_id: str | None = None):
 
         self.project_id = project_id
         self.namespace_id = namespace_id
@@ -220,7 +220,7 @@ async def run_query(self, data) -> list[dict]:
         results, _ = await self._run_query(data)
         return results
 
-    async def _run_query(self, data) -> tuple[list[dict], Optional[str]]:
+    async def _run_query(self, data) -> tuple[list[dict], str| None]:
         results = []
         cursor = None
 
@@ -281,7 +281,7 @@ async def _run_query(self, data) -> tuple[list[dict], Optional[str]]:
         return results, cursor
 
     async def _get_entities_cursor(self, data) -> \
-            tuple[list[Entity], Optional[str]]:
+            tuple[list[Entity], str | None]:
         results, cursor = await self._run_query(data)
         return [Entity(result['entity']) for result in results], cursor
 
@@ -301,7 +301,7 @@ async def get_keys(self, data) -> list[Key]:
         results, _ = await self._run_query(data)
         return [Key(result['entity']['key']) for result in results]
 
-    async def get_entity(self, data) -> Optional[Entity]:
+    async def get_entity(self, data) -> Entity | None:
         """Return an entity object by given query data.
 
         :param data: see the following link for the data format:
@@ -313,18 +313,18 @@ async def get_entity(self, data) -> Optional[Entity]:
         result = await self.get_entities(data)
         return result[0] if result else None
 
-    async def get_key(self, data) -> Optional[Key]:
+    async def get_key(self, data) -> Key | None:
         data['query']['limit'] = 1
         result = await self.get_keys(data)
         return result[0] if result else None
 
-    async def get_entities_by_kind(self, kind: str,
-                                   offset: Optional[int] = None,
-                                   limit: Optional[int] = None,
-                                   cursor: Optional[str] = None) -> Union[
-                                    list[Entity],
-                                    tuple[list[Entity], Optional[str]]
-                                   ]:
+    async def get_entities_by_kind(
+            self,
+            kind: str,
+            offset: int | None = None,
+            limit: int | None = None,
+            cursor: str | None = None
+            ) -> list[Entity] | tuple[list[Entity], str | None]:
         """Returns entities by kind.
 
         When a limit is set, this function returns a list and a cursor.
@@ -344,8 +344,8 @@ async def get_entities_by_kind(self, kind: str,
             return await self._get_entities_cursor(data)
 
     async def get_entities_by_keys(self, keys: Iterable[Key],
-                                   missing: Optional[list[Any]] = None,
-                                   deferred: Optional[list[Key]] = None,
+                                   missing: list[Any] | None = None,
+                                   deferred: list[Key] | None = None,
                                    eventual: bool = False) -> list[Entity]:
         """Returns entity objects for the given keys or an empty list in case
         no entity is found. The order of entities might not be equal to the
@@ -398,9 +398,9 @@ def data():
         return entities
 
     async def get_entity_by_key(self, key: Key,
-                                missing: Optional[list[Any]] = None,
-                                deferred: Optional[list[Key]] = None,
-                                eventual: bool = False) -> Optional[Entity]:
+                                missing: list[Any] | None = None,
+                                deferred: list[Key] | None = None,
+                                eventual: bool = False) -> Entity | None:
         """Returns an entity object for the given key or None in case no
         entity is found.
 
@@ -447,9 +447,9 @@ def __init__(
             self,
             project_id: str,
             service_file: str,
-            session: Optional[aiohttp.ClientSession] = None,
-            scopes: Optional[Iterable[str]] = None,
-            namespace_id: Optional[str] = None):
+            session: aiohttp.ClientSession | None = None,
+            scopes: Iterable[str] | None = None,
+            namespace_id: str | None = None):
 
         scopes = scopes or list(DEFAULT_SCOPES)
         self.project_id = project_id

aiogcd/connector/decoder.py

@@ -25,8 +25,8 @@ def __new__(cls, *args, ks):
         ks += b'=' * (4 - len(ks) % 4)
         ks = base64.b64decode(ks.replace(b'-', b'+').replace(b'_', b'/'))
         decoder.frombytes(ks)
-        decoder._idx = 0
-        decoder.set_end()
+        decoder._idx = 0  # type: ignore
+        decoder.set_end()  # type: ignore
         return decoder
 
     def set_end(self, end=None):
@@ -35,13 +35,13 @@ def set_end(self, end=None):
         self._end = len(self) if end is None else self._idx + end
 
     def __bool__(self):
-        return self._idx < self._end
+        return self._idx < self._end  # type: ignore
 
     def _get8(self):
         if not self:
             raise BufferDecodeError('truncated')
-        c = self[self._idx]
-        self._idx += 1
+        c = self[self._idx]  # type: ignore
+        self._idx += 1  # type: ignore
         return c
 
     def get_var_int32(self):

aiogcd/connector/entity.py

@@ -37,7 +37,7 @@ def __init__(self, entity_res: dict):
         https://cloud.google.com/datastore/docs/reference/rest/v1/Entity
         """
         self.key = Key(entity_res['key'])
-        self._properties = set()
+        self._properties: set[str] = set()
 
         for prop, val in entity_res['properties'].items():
             self._properties.add(prop)

aiogcd/connector/key.py

@@ -4,7 +4,6 @@
     Author: Jeroen van der Heijden <jeroen@cesbit.com>
 """
 import base64
-from typing import Optional
 from .buffer import Buffer
 from .buffer import BufferDecodeError
 from .path import Path
@@ -40,11 +39,12 @@ class Key:
             Key(path=Path(...), project_id="my-project-id")
 """
     _ks = None
+    path: Path
 
-    def __init__(self, *args, ks: Optional[str] = None,
-                 path: Optional[Path] = None,
-                 project_id: Optional[str] = None,
-                 namespace_id: Optional[str] = None):
+    def __init__(self, *args, ks: str | None = None,
+                 path: Path | None = None,
+                 project_id: str | None = None,
+                 namespace_id: str | None = None):
         if len(args) == 1 and isinstance(args[0], dict):
             assert ks is None and path is None and project_id is None, \
                 self.KEY_INIT_MSG
@@ -96,16 +96,17 @@ def encode(self):
         this method for generating an urlsafe key string.
         """
         buffer = Buffer()
-        buffer.add_var_int32(106)
+        buffer.add_var_int32(106)  # type: ignore
 
         # The project id in a key string is prefixed with s~
-        buffer.add_prefixed_string('s~{}'.format(self.project_id))
+        buffer.add_prefixed_string(  # type: ignore
+            's~{}'.format(self.project_id))
 
         self.path.encode(buffer)
 
         if self.namespace_id:
-            buffer.add_var_int32(162)
-            buffer.add_prefixed_string(self.namespace_id)
+            buffer.add_var_int32(162)  # type: ignore
+            buffer.add_prefixed_string(self.namespace_id)  # type: ignore
 
         return buffer
 
@@ -134,54 +135,55 @@ def kind(self):
         return self.path[-1].kind
 
     @property
-    def id(self):
+    def id(self) -> int | str:
         """Shortcut for .path[-1].id"""
         return self.path[-1].id
 
     @staticmethod
-    def _extract_id_or_name(pair):
+    def _extract_id_or_name(pair) -> int | str:
         """Used on __init__."""
         if 'id' in pair:
             return int(pair['id'])
 
         if 'name' in pair:
             return pair['name']
 
-        return None
+        return ''
 
     @staticmethod
-    def _deserialize_ks(ks: str):
+    def _deserialize_ks(ks: str) -> tuple[str | None, str | None, Path]:
         """Returns a tuple with the project_id, namespace_id and Path
         from a key string."""
 
         decoder = Decoder(ks=ks)
-        project_id = None
-        namespace_id = None
-        path = None
+        project_id: str | None = None
+        namespace_id: str | None = None
+        path: Path | None = None
 
         while decoder:
-            tt = decoder.get_var_int32()
+            tt = decoder.get_var_int32()  # type: ignore
 
             if tt == 106:
                 # The project id in a key string is prefixed with s~ which is
                 # not part of the real project id.
-                project_id = decoder.get_prefixed_string()[2:]
+                project_id = decoder.get_prefixed_string()[2:]  # type: ignore
                 continue
 
             if tt == 114:
-                sz = decoder.get_var_int32()
-                decoder.set_end(sz)
+                sz = decoder.get_var_int32()  # type: ignore
+                decoder.set_end(sz)  # type: ignore
                 path = path_from_decoder(decoder)
-                decoder.set_end()
+                decoder.set_end()  # type: ignore
                 continue
 
             if tt == 162:
-                namespace_id = decoder.get_prefixed_string()
+                namespace_id = decoder.get_prefixed_string()  # type: ignore
                 continue
 
             if tt == 0:
                 raise BufferDecodeError('corrupt')
 
+        assert path
         return project_id, namespace_id, path
 
     def get_parent(self):

aiogcd/connector/path.py

@@ -3,18 +3,17 @@
 Created on: May 19, 2017
     Author: Jeroen van der Heijden <jeroen@cesbit.com>
 """
-from typing import Iterable, Union
+from typing import Iterable
 from .pathelement import PathElement
 from .pathelement import path_element_from_decoder
 from .buffer import BufferDecodeError
 
 
 class Path:
 
-    def __init__(self, pairs: Union[
-                Iterable[PathElement],
-                Iterable[tuple[str, Union[int, str]]]]):
-        self._path: tuple[PathElement] = tuple(
+    def __init__(self, pairs: Iterable[PathElement] | Iterable[
+                 tuple[str, int | str]]):
+        self._path: tuple[PathElement, ...] = tuple(
             pe if isinstance(pe, PathElement) else PathElement(*pe)
             for pe in pairs)
 
@@ -43,7 +42,7 @@ def byte_size(self) -> int:
             n += path_element.byte_size
         return n
 
-    def get_as_tuple(self) -> tuple[tuple[str, Union[str, int]], ...]:
+    def get_as_tuple(self) -> tuple[tuple[str, str | int], ...]:
         """Returns a tuple of pairs (tuples) representing the key path of an
         entity. Useful for composing entities with a specific ancestor."""
         return tuple((pe.kind, pe.id) for pe in self._path)

aiogcd/connector/service_account_token.py

@@ -72,6 +72,8 @@ async def ensure_token(self):
 
         else:
             now = datetime.datetime.now()
+            assert self.access_token_acquired_at is not None
+            assert self.access_token_duration is not None
             delta = (now - self.access_token_acquired_at).total_seconds()
             if delta > self.access_token_duration / 2:
                 self.acquiring = asyncio.ensure_future(
@@ -93,6 +95,7 @@ async def _acquire_access_token(self):
         return True
 
     async def _acquire_token(self):
+        assert self.service_data is not None
         assertion = self._generate_assertion()
         url = self.service_data['token_uri']
 
@@ -115,6 +118,7 @@ async def _acquire_token(self):
         return json
 
     def _generate_assertion(self):
+        assert self.service_data is not None
         payload = self._make_gcloud_oauth_body(
         )
 
@@ -127,6 +131,7 @@ def _generate_assertion(self):
         return jwt_token
 
     def _make_gcloud_oauth_body(self):
+        assert self.service_data is not None
         uri = self.service_data['token_uri']
         client_email = self.service_data['client_email']
 

aiogcd/connector/utils.py

@@ -34,7 +34,7 @@ def value_to_dict(val):
                 .decode('utf-8')}
         return {'stringValue': val}
     if isinstance(val, Key):
-        return {'keyValue': val.get_dict()}
+        return {'keyValue': val.get_dict()}  # type: ignore
     if isinstance(val, list):
         return {'arrayValue': {'values': [value_to_dict(v) for v in val]}}
     if isinstance(val, TimestampValue):