dataclass for link
All checks were successful
continuous-integration/drone Build is passing

This commit is contained in:
Roy Olav Purser 2022-02-12 13:51:17 +01:00
parent fb48a2181a
commit c969a9a1ce
Signed by: roypur
GPG Key ID: E14D26A036F21656
2 changed files with 29 additions and 25 deletions

View File

@ -4,7 +4,8 @@ import os
import pprint import pprint
import re import re
import sys import sys
from typing import Literal, Optional, cast import dataclasses
from typing import Literal, cast
import pydantic import pydantic
@ -41,20 +42,26 @@ stream_server = os.environ.get("STREAM_SERVER")
proxy_server = os.environ.get("PROXY_SERVER") proxy_server = os.environ.get("PROXY_SERVER")
@dataclasses.dataclass
class LinkWithType:
upstream: str
ctype: str
class ProxyCreateLink(pydantic.BaseModel): class ProxyCreateLink(pydantic.BaseModel):
upstream: Optional[pydantic.HttpUrl] upstream: None | pydantic.HttpUrl
ctype: Optional[pydantic.StrictStr] ctype: None | pydantic.StrictStr
region: Optional[pydantic.StrictStr] region: None | pydantic.StrictStr
class ProxyRequest(pydantic.BaseModel): class ProxyRequest(pydantic.BaseModel):
action: Literal["create-urls", "read-config"] action: Literal["create-urls", "read-config"]
urls: Optional[list[ProxyCreateLink]] urls: None | list[ProxyCreateLink]
class ProxyResponse(pydantic.BaseModel): class ProxyResponse(pydantic.BaseModel):
action: Literal["create-urls", "read-config"] action: Literal["create-urls", "read-config"]
urls: Optional[list[pydantic.HttpUrl]] urls: None | list[pydantic.HttpUrl]
class ProxyElem: class ProxyElem:
@ -69,36 +76,30 @@ class ProxyElem:
def __repr__(self): def __repr__(self):
return f"<region=<{str(self.region)}> proxy=<{str(self.proxy)}>>" return f"<region=<{str(self.region)}> proxy=<{str(self.proxy)}>>"
async def proxy_url(self, urls): async def proxy_url(self, urls: list[LinkWithType]):
clean_urls = [] clean_urls = []
link_requests: list[ProxyCreateLink] = [] link_requests: list[ProxyCreateLink] = []
for url in urls: for url in urls:
if isinstance(url, tuple): clean_urls.append(url.upstream)
clean_urls.append(url[0])
else:
clean_urls.append(url)
if not isinstance(proxy_server, str): if not isinstance(proxy_server, str):
return clean_urls return clean_urls
try: try:
for url in urls: for url in urls:
if isinstance(url, tuple): link_requests.append(
link_requests.append( ProxyCreateLink.parse_obj(
ProxyCreateLink.parse_obj( {
{"upstream": url[0], "ctype": url[1], "region": self.region} "upstream": url.upstream,
) "ctype": url.ctype,
) "region": self.region,
else: }
link_requests.append(
ProxyCreateLink.parse_obj(
{"upstream": url, "region": self.region}
)
) )
)
except pydantic.ValidationError as e: except pydantic.ValidationError as e:
logger.info(e) logger.info(e)
return clean_urls return clean_urls
response_data: ProxyRequest request_data: ProxyRequest
response_data: ProxyResponse response_data: ProxyResponse
try: try:
@ -118,7 +119,7 @@ class ProxyElem:
except (aiohttp.ClientError, pydantic.ValidationError) as e: except (aiohttp.ClientError, pydantic.ValidationError) as e:
logger.info(e) logger.info(e)
else: else:
ret_data = [] ret_data: list[None | str] = []
if response_data.urls is not None: if response_data.urls is not None:
for src, dst in zip(clean_urls, response_data.urls): for src, dst in zip(clean_urls, response_data.urls):
if isinstance(src, str): if isinstance(src, str):

View File

@ -134,7 +134,10 @@ class MainHandler(tornado.web.RequestHandler):
else: else:
proxied = await handler.proxy.proxy_url( proxied = await handler.proxy.proxy_url(
[ [
(provider_data.upstream(), provider_data.proxy_ctype()), config.LinkWithType(
upstream=provider_data.upstream(),
ctype=provider_data.proxy_ctype(),
),
provider_data.thumbnail(), provider_data.thumbnail(),
] ]
) )