Merge pull request #2 from chriskuehl/fix-type-checking

Fix types

Author: chriskuehl

poetry.lock

@@ -5,6 +5,7 @@
 import mimetypes
 import os.path
 import re
+import typing
 
 import fluffy_code.code
 import fluffy_code.prebuilt_styles
@@ -60,7 +61,11 @@
 
 class CacheControlHeaderMiddleware(BaseHTTPMiddleware):
 
-    async def dispatch(self, request, call_next):
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: typing.Callable[[Request], typing.Awaitable[Response]],
+    ) -> Response:
         response = await call_next(request)
         # TODO: There should be a better way to do this...
         response.headers['Cache-Control'] = 'no-cache'
@@ -267,9 +272,9 @@ async def package_file_archive_path(request: Request) -> Response:
         ('Mode', Markup(f'<tt>{entry.mode}</tt>')),
     )
 
-    def _transfer_raw():
+    def _transfer_raw() -> Response:
         """Return the file verbatim."""
-        async def transfer_file():
+        async def transfer_file() -> typing.AsyncIterator[bytes]:
             async with package.open_from_archive(archive_path) as f:
                 data = None
                 while data is None or len(data) > 0:
@@ -306,10 +311,15 @@ async def transfer_file():
             style_config = fluffy_code.prebuilt_styles.default_style()
 
             try:
-                lexer = pygments.lexers.guess_lexer_for_filename(
-                    archive_path,
-                    text,
-                    stripnl=False,
+                # TODO: Remove this cast once
+                # https://github.com/python/typeshed/pull/8777 is merged.
+                lexer = typing.cast(
+                    pygments.lexer.Lexer,
+                    pygments.lexers.guess_lexer_for_filename(
+                        archive_path,
+                        text,
+                        stripnl=False,
+                    ),
                 )
             except pygments.lexers.ClassNotFound:
                 lexer = pygments.lexers.special.TextLexer(stripnl=False)

pypi_browser/app.py

@@ -9,6 +9,7 @@
 import typing
 import zipfile
 from dataclasses import dataclass
+from types import TracebackType
 
 
 def pep426_normalize(package_name: str) -> str:
@@ -63,23 +64,25 @@ def _package_entries_from_tarball(path: str) -> typing.Set[PackageEntry]:
         }
 
 
-ArchiveFile = typing.Union[zipfile.ZipExtFile]
-
-
 class AsyncArchiveFile:
 
-    file_: ArchiveFile
+    file_: typing.IO[bytes]
 
-    def __init__(self, file_: ArchiveFile) -> None:
+    def __init__(self, file_: typing.IO[bytes]) -> None:
         self.file_ = file_
 
     async def __aenter__(self) -> 'AsyncArchiveFile':
         return self
 
-    async def __aexit__(self, exc_t, exc_v, exc_tb) -> None:
+    async def __aexit__(
+        self,
+        exc_t: typing.Optional[typing.Type[BaseException]],
+        exc_v: typing.Optional[BaseException],
+        exc_tb: typing.Optional[TracebackType],
+    ) -> None:
         await asyncio.to_thread(self.file_.close)
 
-    async def read(self, n_bytes: typing.Optional[int] = None) -> bytes:
+    async def read(self, n_bytes: int = -1) -> bytes:
         return await asyncio.to_thread(self.file_.read, n_bytes)
 
 
@@ -123,21 +126,22 @@ async def entries(self) -> typing.Set[PackageEntry]:
             raise AssertionError(self.package_format)
 
     @contextlib.asynccontextmanager
-    async def open_from_archive(self, path: str) -> str:
+    async def open_from_archive(self, path: str) -> typing.AsyncIterator[AsyncArchiveFile]:
         if self.package_format is PackageFormat.ZIPFILE:
             zf = await asyncio.to_thread(zipfile.ZipFile, self.path)
-            archive_file = await asyncio.to_thread(zf.open, path)
             try:
-                async with AsyncArchiveFile(archive_file) as zip_archive_file:
-                    yield zip_archive_file
+                zip_archive_file = await asyncio.to_thread(zf.open, path)
+                async with AsyncArchiveFile(zip_archive_file) as wrapped:
+                    yield wrapped
             finally:
                 await asyncio.to_thread(zf.close)
         elif self.package_format is PackageFormat.TARBALL:
             tf = await asyncio.to_thread(tarfile.open, self.path)
-            archive_file = await asyncio.to_thread(tf.extractfile, path)
             try:
-                async with AsyncArchiveFile(archive_file) as zip_archive_file:
-                    yield zip_archive_file
+                tar_archive_file = await asyncio.to_thread(tf.extractfile, path)
+                assert tar_archive_file is not None, path
+                async with AsyncArchiveFile(tar_archive_file) as wrapped:
+                    yield wrapped
             finally:
                 await asyncio.to_thread(tf.close)
         else:

pypi_browser/packaging.py

@@ -20,7 +20,11 @@ class PackageDoesNotExist(Exception):
     pass
 
 
-async def package_metadata(config: PyPIConfig, client: httpx.AsyncClient, package: str) -> typing.Dict:
+async def package_metadata(
+    config: PyPIConfig,
+    client: httpx.AsyncClient,
+    package: str,
+) -> typing.Dict[typing.Any, typing.Any]:
     resp = await client.get(f'{config.pypi_url}/pypi/{package}/json')
     if resp.status_code == 404:
         raise PackageDoesNotExist(package)
@@ -54,16 +58,17 @@ def _storage_path(config: PyPIConfig, package: str, filename: str) -> str:
 
 
 @contextlib.asynccontextmanager
-async def _atomic_file(path: str, mode: str = 'w') -> typing.Any:
-    async with aiofiles.tempfile.NamedTemporaryFile(mode, dir=os.path.dirname(path), delete=False) as f:
+async def _atomic_file(path: str) -> typing.AsyncIterator[aiofiles.threadpool.binary.AsyncBufferedIOBase]:
+    async with aiofiles.tempfile.NamedTemporaryFile('wb', dir=os.path.dirname(path), delete=False) as f:
+        tmp_path = typing.cast(str, f.name)
         try:
             yield f
         except BaseException:
-            await aiofiles.os.remove(f.name)
+            await aiofiles.os.remove(tmp_path)
             raise
         else:
             # This is atomic since the temporary file was created in the same directory.
-            await aiofiles.os.rename(f.name, path)
+            await aiofiles.os.rename(tmp_path, path)
 
 
 async def downloaded_file_path(config: PyPIConfig, package: str, filename: str) -> str:
@@ -91,7 +96,7 @@ async def downloaded_file_path(config: PyPIConfig, package: str, filename: str)
 
         await aiofiles.os.makedirs(os.path.dirname(stored_path), exist_ok=True)
 
-        async with _atomic_file(stored_path, 'wb') as f:
+        async with _atomic_file(stored_path) as f:
             async with client.stream('GET', url) as resp:
                 resp.raise_for_status()
                 async for chunk in resp.aiter_bytes():

pypi_browser/pypi.py

@@ -27,6 +27,7 @@ pytest = "^7.1.3"
 mypy = "^0.971"
 coverage = "^6.4.4"
 types-aiofiles = "^22.1.0"
+types-Pygments = "^2.13.0"
 
 [build-system]
 requires = ["poetry-core"]