implemented some pydandic models and fixed type errors

.gitignore

@@ -8,3 +8,4 @@
 # python
 __pycache__/
 .pytest_cache/
+.mypy_cache/

flake.nix

@@ -23,7 +23,8 @@
     packages = forAllSystems (system: {
       default = poetry2nix.${system}.mkPoetryApplication {
         projectDir = self;
-        checkPhase = "pytest";
+        checkPhase = "mypy . && pytest";
+        preferWheels = true; # avoid long build process for `mypy`
       };
       vm = self.nixosConfigurations.vm.config.microvm.declaredRunner;
     });

poetry.lock

@@ -504,6 +504,20 @@ anyio = ">=3.4.0,<5"
 [package.extras]
 full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart (>=0.0.7)", "pyyaml"]
 
+[[package]]
+name = "types-requests"
+version = "2.32.0.20240712"
+description = "Typing stubs for requests"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "types-requests-2.32.0.20240712.tar.gz", hash = "sha256:90c079ff05e549f6bf50e02e910210b98b8ff1ebdd18e19c873cd237737c1358"},
+    {file = "types_requests-2.32.0.20240712-py3-none-any.whl", hash = "sha256:f754283e152c752e46e70942fa2a146b5bc70393522257bb85bd1ef7e019dcc3"},
+]
+
+[package.dependencies]
+urllib3 = ">=2"
+
 [[package]]
 name = "typing-extensions"
 version = "4.12.2"
@@ -553,4 +567,4 @@ standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)",
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "a1f2445b7578ac1606a1968f771da2d0c31e859298a4d8815cfbfcf94d06ac57"
+content-hash = "8c95de1de17ced501a0460661e22c77839a3afca88512e07ad95344e719679a5"

pyproject.toml

@@ -23,6 +23,7 @@ testdata = "testdata.main:main"
 pytest = "^8.3.2"
 requests = "^2.32.3"
 mypy = "^1.11.1"
+types-requests = "^2.32.0.20240712"
 
 [build-system]
 requires = ["poetry-core"]

testdata/api.py

@@ -1,14 +1,14 @@
 from fastapi import FastAPI, HTTPException, status
 from fastapi.responses import StreamingResponse
+from pydantic import ValidationError
 
-from .utils import convert_to_bytes, load_database, save_database, generate_data
+from .utils import load_database, save_database, generate_data
+from .custom_types import TestDataBody
 
-
-def create_api(api_keys: {str}, max_size: int, max_data: int, database: str, buffer_size: int):
+def create_api(api_keys: set[str], max_size: int, max_data: int, database: str, buffer_size: int):
 
     api = FastAPI(docs_url=None, redoc_url=None)
 
-
     class MaxSizePerRequestError(Exception):
         pass
 
@@ -20,41 +20,41 @@ def create_api(api_keys: {str}, max_size: int, max_data: int, database: str, buf
     @api.get('/')
     async def test_data(api_key: str, size: str) -> StreamingResponse:
         try:
-            if api_key not in api_keys:
+            body = TestDataBody(api_key=api_key, size=size)  # type: ignore
+        except ValidationError as err:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail='Invalid Format.'
+            ) from err
+
+        try:
+            if body.api_key not in api_keys:
                 raise HTTPException(
                     status_code=status.HTTP_401_UNAUTHORIZED,
                     detail='Invalid API Key.'
                 )
 
-            try:
-                size = convert_to_bytes(size)
-            except ValueError as err:
-                raise HTTPException(
-                    status_code=status.HTTP_400_BAD_REQUEST,
-                    detail='Invalid format format for size.'
-                ) from err
-
-            if size < 0:
+            if body.size < 0:
                 raise MinSizePerRequestError
-            elif max_size < size:
+            elif max_size < body.size:
                 raise MaxSizePerRequestError
 
             db = load_database(database)
-            if max_data <= db['data-used'] + size:
+            if max_data <= db['data-used'] + body.size:
                 raise HTTPException(
                     status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                     detail='Service not available.'
                 )
-            db['data-used'] += size
+            db['data-used'] += body.size
 
             save_database(database, db)
 
             return StreamingResponse(
                 status_code=status.HTTP_200_OK,
-                content=generate_data(size, buffer_size),
+                content=generate_data(body.size, buffer_size),
                 media_type='application/octet-stream',
                 headers={
-                    'Content-Length': str(size)
+                    'Content-Length': str(body.size)
                 }
             )
 

testdata/custom_types.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel, ConfigDict, field_validator
+
+from .utils import convert_to_bytes
+
+# class Config(BaseModel):
+#     host: str
+#     port: int
+#     buffer_size: int = 4 * 1024  # 4KB
+#     max_size: int = 2 * 1024 * 1024 * 1024  # 2GB
+#     max_data: int = 0  # unlimited
+#     api_keys = set[str]
+#     database = str
+
+#     model_config = ConfigDict(extra='allow')
+
+#     @field_validator('buffer_size', 'max_size', 'max_data')
+#     @classmethod
+#     def convert_size(cls, value: int | str) -> int:
+#         return convert_to_bytes(value)
+
+
+# TODO
+# class DataBase(BaseModel):
+#     model_config = ConfigDict(extra='forbid')
+
+
+class TestDataBody(BaseModel):
+    api_key: str
+    size: int
+
+    model_config = ConfigDict(extra='forbid')
+
+    @field_validator('size', mode='before')
+    @classmethod
+    def convert_size(cls, value: str) -> int:
+        return convert_to_bytes(value)

testdata/main.py

@@ -9,13 +9,22 @@ import uvicorn
 from .utils import convert_to_bytes, save_database
 from .api import create_api
 
+
 # Setup Parser
+def parse_cli_arguments(argv: list[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser()
-parser.add_argument('-c', '--config', type=argparse.FileType('r'),
-                    default='./config.json', help='Path to config file in JSON format.')
+    parser.add_argument(
+        '-c',
+        '--config',
+        type = argparse.FileType('r'),
+        default = './config.json',
+        help = 'Path to config file in JSON format.'
+    )
+
+    return parser.parse_args(argv)
 
 
-def run(host: str, port: int, api_keys: {str}, max_size: int, max_data: int, database: str, buffer_size: int):
+def run(host: str, port: int, api_keys: set[str], max_size: int, max_data: int, database: str, buffer_size: int):
     if not exists(database) or os.stat(database).st_size == 0:
         save_database(database, {'data-used': 0})
 
@@ -25,8 +34,9 @@ def run(host: str, port: int, api_keys: {str}, max_size: int, max_data: int, dat
 
 
 def main():
-    args = parser.parse_args(sys.argv[1:])
+    args = parse_cli_arguments(sys.argv[1:])
     config = json.load(args.config)
+
     host = config['host']
     port = config['port']
     buffer_size = convert_to_bytes(config['buffer-size'])

testdata/utils.py

@@ -1,10 +1,15 @@
 import json
 import asyncio
+from typing import AsyncGenerator
+
 
 def convert_to_bytes(size: int | str) -> int:
+    if isinstance(size, int):
+        return size
+
     try:
         return int(size)
-    except ValueError:  # treat as string
+    except ValueError:
         units = {
             'TB': 1000 ** 4, 'TiB': 1024 ** 4,
             'GB': 1000 ** 3, 'GiB': 1024 ** 3,
@@ -16,12 +21,11 @@ def convert_to_bytes(size: int | str) -> int:
         for unit in units:
             if size.endswith(unit):
                 return int(float(size.removesuffix(unit)) * units[unit])
-                break
-
-        raise ValueError
+        else:
+            raise ValueError('Invalid format. Expected integer or float ending with a data unit (B, KB, MiB,...).')
 
 
-async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
+async def generate_data(size: int, buffer_size: int = 4 * 1024) -> AsyncGenerator[bytes, None]:
     size_left = size
 
     # https://github.com/tiangolo/fastapi/issues/5183
@@ -39,11 +43,6 @@ async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes:
         raise GeneratorExit
 
 
-def check_policies(ip: str) -> None:
-    network = ipaddress.ip_network(ip)
-    print(network)
-
-
 def load_database(path: str) -> dict:
     with open(path, 'r') as file:
         return json.load(file)