implemented some pydandic models and fixed type errors
This commit is contained in:
parent
93c51a00a0
commit
3f8323d297
1
.gitignore
vendored
1
.gitignore
vendored
@ -8,3 +8,4 @@
|
||||
# python
|
||||
__pycache__/
|
||||
.pytest_cache/
|
||||
.mypy_cache/
|
||||
|
@ -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;
|
||||
});
|
||||
|
16
poetry.lock
generated
16
poetry.lock
generated
@ -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"
|
||||
|
@ -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"]
|
||||
|
38
testdata/api.py
vendored
38
testdata/api.py
vendored
@ -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)
|
||||
}
|
||||
)
|
||||
|
||||
|
36
testdata/custom_types.py
vendored
Normal file
36
testdata/custom_types.py
vendored
Normal file
@ -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)
|
20
testdata/main.py
vendored
20
testdata/main.py
vendored
@ -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'])
|
||||
|
19
testdata/utils.py
vendored
19
testdata/utils.py
vendored
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user