From 3f8323d29726a5785d1895e37678b148ddb94620 Mon Sep 17 00:00:00 2001 From: Kristian Krsnik Date: Sat, 3 Aug 2024 00:52:27 +0200 Subject: [PATCH] implemented some pydandic models and fixed type errors --- .gitignore | 1 + flake.nix | 3 ++- poetry.lock | 16 +++++++++++++++- pyproject.toml | 1 + testdata/api.py | 38 +++++++++++++++++++------------------- testdata/custom_types.py | 36 ++++++++++++++++++++++++++++++++++++ testdata/main.py | 20 +++++++++++++++----- testdata/utils.py | 19 +++++++++---------- 8 files changed, 98 insertions(+), 36 deletions(-) create mode 100644 testdata/custom_types.py diff --git a/.gitignore b/.gitignore index 950514e..f4aaf52 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ # python __pycache__/ .pytest_cache/ +.mypy_cache/ diff --git a/flake.nix b/flake.nix index 8e7b007..9e98fa6 100644 --- a/flake.nix +++ b/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; }); diff --git a/poetry.lock b/poetry.lock index b34c9b3..ae554e1 100644 --- a/poetry.lock +++ b/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" diff --git a/pyproject.toml b/pyproject.toml index 3cfeb43..a9fbe8e 100644 --- a/pyproject.toml +++ b/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"] diff --git a/testdata/api.py b/testdata/api.py index 4863d73..f957e0b 100644 --- a/testdata/api.py +++ b/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) } ) diff --git a/testdata/custom_types.py b/testdata/custom_types.py new file mode 100644 index 0000000..142624c --- /dev/null +++ b/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) diff --git a/testdata/main.py b/testdata/main.py index f03606c..62b57f8 100644 --- a/testdata/main.py +++ b/testdata/main.py @@ -9,13 +9,22 @@ import uvicorn from .utils import convert_to_bytes, save_database from .api import create_api + # Setup Parser -parser = argparse.ArgumentParser() -parser.add_argument('-c', '--config', type=argparse.FileType('r'), - default='./config.json', help='Path to config file in JSON format.') +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.' + ) + + 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']) diff --git a/testdata/utils.py b/testdata/utils.py index 8a9a09d..bde8924 100644 --- a/testdata/utils.py +++ b/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)