Add config options

Author: chriskuehl

pypi_browser/app.py

@@ -9,6 +9,7 @@
 import fluffy_code.prebuilt_styles
 import packaging.version
 import pygments.lexers.special
+import starlette.config
 from identify import identify
 from markupsafe import Markup
 from starlette.applications import Starlette
@@ -25,6 +26,7 @@
 import pypi_browser.packaging
 from pypi_browser import pypi
 
+
 PACKAGE_TYPE_NOT_SUPPORTED_ERROR = (
     'Sorry, this package type is not yet supported (only .zip and .whl supported currently).'
 )
@@ -65,6 +67,12 @@ async def dispatch(self, request, call_next):
         return response
 
 
+config = starlette.config.Config()
+pypi_config = pypi.PyPIConfig(
+    cache_path=config('PYPI_BROWSER_PACKAGE_CACHE_PATH', default='/tmp'),
+    pypi_url=config('PYPI_BROWSER_PYPI_URL', default='https://pypi.org'),
+)
+
 app = Starlette(
     debug=os.environ.get('PYPI_BROWSER_DEBUG') == '1',
     middleware=[
@@ -110,7 +118,7 @@ async def package(request: Request) -> Response:
         return RedirectResponse(request.url_for('package', package=normalized_package_name))
 
     try:
-        version_to_files = await pypi.files_for_package(package_name)
+        version_to_files = await pypi.files_for_package(pypi_config, package_name)
     except pypi.PackageDoesNotExist:
         return PlainTextResponse(
             f'Package {package_name!r} does not exist on PyPI.',
@@ -149,7 +157,7 @@ async def package_file(request: Request) -> Response:
         )
 
     try:
-        archive = await pypi.downloaded_file_path(package_name, file_name)
+        archive = await pypi.downloaded_file_path(pypi_config, package_name, file_name)
     except pypi.PackageDoesNotExist:
         return PlainTextResponse(
             f'Package {package_name!r} does not exist on PyPI.',
@@ -223,7 +231,7 @@ async def package_file_archive_path(request: Request) -> Response:
         )
 
     try:
-        archive = await pypi.downloaded_file_path(package_name, file_name)
+        archive = await pypi.downloaded_file_path(pypi_config, package_name, file_name)
     except pypi.PackageDoesNotExist:
         return PlainTextResponse(
             f'Package {package_name!r} does not exist on PyPI.',

pypi_browser/pypi.py

@@ -1,5 +1,6 @@
 import base64
 import contextlib
+import dataclasses
 import itertools
 import os.path
 import typing
@@ -8,26 +9,27 @@
 import httpx
 
 
-# TODO: make configurable
-PYPI = 'https://pypi.org'
-STORAGE_DIR = '/home/ckuehl/tmp/pypi-view'
+@dataclasses.dataclass(frozen=True)
+class PyPIConfig:
+    cache_path: str
+    pypi_url: str
 
 
 class PackageDoesNotExist(Exception):
     pass
 
 
-async def package_metadata(client: httpx.AsyncClient, package: str) -> typing.Dict:
-    resp = await client.get(f'{PYPI}/pypi/{package}/json')
+async def package_metadata(config: PyPIConfig, client: httpx.AsyncClient, package: str) -> typing.Dict:
+    resp = await client.get(f'{config.pypi_url}/pypi/{package}/json')
     if resp.status_code == 404:
         raise PackageDoesNotExist(package)
     resp.raise_for_status()
     return resp.json()
 
 
-async def files_for_package(package: str) -> typing.Dict[str, typing.Set[str]]:
+async def files_for_package(config: PyPIConfig, package: str) -> typing.Dict[str, typing.Set[str]]:
     async with httpx.AsyncClient() as client:
-        metadata = await package_metadata(client, package)
+        metadata = await package_metadata(config, client, package)
 
     return {
         version: {file_['filename'] for file_ in files}
@@ -40,9 +42,9 @@ class CannotFindFileError(Exception):
     pass
 
 
-def _storage_path(package: str, filename: str) -> str:
+def _storage_path(config: PyPIConfig, package: str, filename: str) -> str:
     return os.path.join(
-        STORAGE_DIR,
+        config.cache_path,
         # Base64-encoding the names to calculate the storage path just to be
         # extra sure to avoid any path traversal vulnerabilities.
         base64.urlsafe_b64encode(package.encode('utf8')).decode('ascii'),
@@ -63,18 +65,18 @@ async def _atomic_file(path: str, mode: str = 'w') -> typing.Any:
             await aiofiles.os.rename(f.name, path)
 
 
-async def downloaded_file_path(package: str, filename: str) -> str:
+async def downloaded_file_path(config: PyPIConfig, package: str, filename: str) -> str:
     """Return path on filesystem to downloaded PyPI file.
 
     May be instant if the file is already cached; otherwise it will download
     it and may take a while.
     """
-    stored_path = _storage_path(package, filename)
+    stored_path = _storage_path(config, package, filename)
     if await aiofiles.os.path.exists(stored_path):
         return stored_path
 
     async with httpx.AsyncClient() as client:
-        metadata = await package_metadata(client, package)
+        metadata = await package_metadata(config, client, package)
 
         # Parsing versions from non-wheel Python packages isn't perfectly
         # reliable, so just search through all releases until we find a