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
|
# python
|
||||||
__pycache__/
|
__pycache__/
|
||||||
.pytest_cache/
|
.pytest_cache/
|
||||||
|
.mypy_cache/
|
||||||
|
@ -23,7 +23,8 @@
|
|||||||
packages = forAllSystems (system: {
|
packages = forAllSystems (system: {
|
||||||
default = poetry2nix.${system}.mkPoetryApplication {
|
default = poetry2nix.${system}.mkPoetryApplication {
|
||||||
projectDir = self;
|
projectDir = self;
|
||||||
checkPhase = "pytest";
|
checkPhase = "mypy . && pytest";
|
||||||
|
preferWheels = true; # avoid long build process for `mypy`
|
||||||
};
|
};
|
||||||
vm = self.nixosConfigurations.vm.config.microvm.declaredRunner;
|
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]
|
[package.extras]
|
||||||
full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart (>=0.0.7)", "pyyaml"]
|
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]]
|
[[package]]
|
||||||
name = "typing-extensions"
|
name = "typing-extensions"
|
||||||
version = "4.12.2"
|
version = "4.12.2"
|
||||||
@ -553,4 +567,4 @@ standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)",
|
|||||||
[metadata]
|
[metadata]
|
||||||
lock-version = "2.0"
|
lock-version = "2.0"
|
||||||
python-versions = "^3.11"
|
python-versions = "^3.11"
|
||||||
content-hash = "a1f2445b7578ac1606a1968f771da2d0c31e859298a4d8815cfbfcf94d06ac57"
|
content-hash = "8c95de1de17ced501a0460661e22c77839a3afca88512e07ad95344e719679a5"
|
||||||
|
@ -23,6 +23,7 @@ testdata = "testdata.main:main"
|
|||||||
pytest = "^8.3.2"
|
pytest = "^8.3.2"
|
||||||
requests = "^2.32.3"
|
requests = "^2.32.3"
|
||||||
mypy = "^1.11.1"
|
mypy = "^1.11.1"
|
||||||
|
types-requests = "^2.32.0.20240712"
|
||||||
|
|
||||||
[build-system]
|
[build-system]
|
||||||
requires = ["poetry-core"]
|
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 import FastAPI, HTTPException, status
|
||||||
from fastapi.responses import StreamingResponse
|
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: set[str], max_size: int, max_data: int, database: str, buffer_size: int):
|
||||||
def create_api(api_keys: {str}, max_size: int, max_data: int, database: str, buffer_size: int):
|
|
||||||
|
|
||||||
api = FastAPI(docs_url=None, redoc_url=None)
|
api = FastAPI(docs_url=None, redoc_url=None)
|
||||||
|
|
||||||
|
|
||||||
class MaxSizePerRequestError(Exception):
|
class MaxSizePerRequestError(Exception):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@ -20,41 +20,41 @@ def create_api(api_keys: {str}, max_size: int, max_data: int, database: str, buf
|
|||||||
@api.get('/')
|
@api.get('/')
|
||||||
async def test_data(api_key: str, size: str) -> StreamingResponse:
|
async def test_data(api_key: str, size: str) -> StreamingResponse:
|
||||||
try:
|
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(
|
raise HTTPException(
|
||||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||||
detail='Invalid API Key.'
|
detail='Invalid API Key.'
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
if body.size < 0:
|
||||||
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:
|
|
||||||
raise MinSizePerRequestError
|
raise MinSizePerRequestError
|
||||||
elif max_size < size:
|
elif max_size < body.size:
|
||||||
raise MaxSizePerRequestError
|
raise MaxSizePerRequestError
|
||||||
|
|
||||||
db = load_database(database)
|
db = load_database(database)
|
||||||
if max_data <= db['data-used'] + size:
|
if max_data <= db['data-used'] + body.size:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||||
detail='Service not available.'
|
detail='Service not available.'
|
||||||
)
|
)
|
||||||
db['data-used'] += size
|
db['data-used'] += body.size
|
||||||
|
|
||||||
save_database(database, db)
|
save_database(database, db)
|
||||||
|
|
||||||
return StreamingResponse(
|
return StreamingResponse(
|
||||||
status_code=status.HTTP_200_OK,
|
status_code=status.HTTP_200_OK,
|
||||||
content=generate_data(size, buffer_size),
|
content=generate_data(body.size, buffer_size),
|
||||||
media_type='application/octet-stream',
|
media_type='application/octet-stream',
|
||||||
headers={
|
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 .utils import convert_to_bytes, save_database
|
||||||
from .api import create_api
|
from .api import create_api
|
||||||
|
|
||||||
|
|
||||||
# Setup Parser
|
# Setup Parser
|
||||||
parser = argparse.ArgumentParser()
|
def parse_cli_arguments(argv: list[str]) -> argparse.Namespace:
|
||||||
parser.add_argument('-c', '--config', type=argparse.FileType('r'),
|
parser = argparse.ArgumentParser()
|
||||||
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:
|
if not exists(database) or os.stat(database).st_size == 0:
|
||||||
save_database(database, {'data-used': 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():
|
def main():
|
||||||
args = parser.parse_args(sys.argv[1:])
|
args = parse_cli_arguments(sys.argv[1:])
|
||||||
config = json.load(args.config)
|
config = json.load(args.config)
|
||||||
|
|
||||||
host = config['host']
|
host = config['host']
|
||||||
port = config['port']
|
port = config['port']
|
||||||
buffer_size = convert_to_bytes(config['buffer-size'])
|
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 json
|
||||||
import asyncio
|
import asyncio
|
||||||
|
from typing import AsyncGenerator
|
||||||
|
|
||||||
|
|
||||||
def convert_to_bytes(size: int | str) -> int:
|
def convert_to_bytes(size: int | str) -> int:
|
||||||
|
if isinstance(size, int):
|
||||||
|
return size
|
||||||
|
|
||||||
try:
|
try:
|
||||||
return int(size)
|
return int(size)
|
||||||
except ValueError: # treat as string
|
except ValueError:
|
||||||
units = {
|
units = {
|
||||||
'TB': 1000 ** 4, 'TiB': 1024 ** 4,
|
'TB': 1000 ** 4, 'TiB': 1024 ** 4,
|
||||||
'GB': 1000 ** 3, 'GiB': 1024 ** 3,
|
'GB': 1000 ** 3, 'GiB': 1024 ** 3,
|
||||||
@ -16,12 +21,11 @@ def convert_to_bytes(size: int | str) -> int:
|
|||||||
for unit in units:
|
for unit in units:
|
||||||
if size.endswith(unit):
|
if size.endswith(unit):
|
||||||
return int(float(size.removesuffix(unit)) * units[unit])
|
return int(float(size.removesuffix(unit)) * units[unit])
|
||||||
break
|
else:
|
||||||
|
raise ValueError('Invalid format. Expected integer or float ending with a data unit (B, KB, MiB,...).')
|
||||||
raise ValueError
|
|
||||||
|
|
||||||
|
|
||||||
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
|
size_left = size
|
||||||
|
|
||||||
# https://github.com/tiangolo/fastapi/issues/5183
|
# 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
|
raise GeneratorExit
|
||||||
|
|
||||||
|
|
||||||
def check_policies(ip: str) -> None:
|
|
||||||
network = ipaddress.ip_network(ip)
|
|
||||||
print(network)
|
|
||||||
|
|
||||||
|
|
||||||
def load_database(path: str) -> dict:
|
def load_database(path: str) -> dict:
|
||||||
with open(path, 'r') as file:
|
with open(path, 'r') as file:
|
||||||
return json.load(file)
|
return json.load(file)
|
||||||
|
Loading…
Reference in New Issue
Block a user