[~] Refactor

This commit is contained in:
Siarhei Siniak 2024-07-06 20:47:12 +03:00
parent 639e755dbd
commit b2afd124a7
3 changed files with 57 additions and 12 deletions

@ -11,6 +11,8 @@ RUN pip3 install tqdm
RUN apt-get install -yy ffmpeg RUN apt-get install -yy ffmpeg
RUN pip3 install celery redis RUN pip3 install celery redis
RUN pip3 install dataclasses-json RUN pip3 install dataclasses-json
RUN pip3 install rpdb
RUN apt-get install -yy netcat-traditional
WORKDIR /app WORKDIR /app

@ -6,16 +6,18 @@ import redis
c = tiktok_config() c = tiktok_config()
kombu_register_json_dataclass()
app = celery.Celery( app = celery.Celery(
__name__, __name__,
broker=c.celery_broker, broker=c.celery_broker,
result_backend=c.celery_result_backend, result_backend=c.celery_result_backend,
accept_content=['json-dataclass', 'json',], accept_content=['json-dataclass'],
task_serializer='json-dataclass', task_serializer='json-dataclass',
result_serializer='json-dataclass',
) )
kombu_register_json_dataclass()
app.autodiscover_tasks(c.celery_imports) app.autodiscover_tasks(c.celery_imports)
redis = dict( redis = dict(

@ -3,15 +3,21 @@ import importlib
import kombu.utils.json import kombu.utils.json
from typing import ( from typing import (
Any, Any,
Optional,
) )
class Task(celery.app.task.Task): class Task(celery.app.task.Task):
def __call__(self, *args, **kwargs): def __call__(self, *args, **kwargs):
res = super().__call__(*args, **kwargs) res = super().__call__(*args, **kwargs)
return self._dumps(res) return res
@classmethod @classmethod
def _loads(self, data_str: str) -> Any: def _loads(
cls,
data_str: Optional[str]=None,
data: Optional[Any]=None,
) -> Any:
if not data_str is None:
data = kombu.utils.json.loads(data_str) data = kombu.utils.json.loads(data_str)
if isinstance(data, dict) and data.get('type') == 'dataclass_json': if isinstance(data, dict) and data.get('type') == 'dataclass_json':
@ -25,20 +31,55 @@ class Task(celery.app.task.Task):
c = getattr(c, current_name) c = getattr(c, current_name)
return c(**data['data']) return c(**data['data'])
else:
if isinstance(data, list):
return [
cls._loads(data=o)
for o in data
]
elif isinstance(data, dict):
return {
k : cls._loads(data=v)
for k, v in data.items()
}
else: else:
return data return data
@classmethod @classmethod
def _dumps(self, data: Any) -> str: def _dumps(
cls,
data: Any,
need_native: Optional[bool]=None,
) -> Any:
if need_native is None:
need_native = False
native = None
if hasattr(data, 'to_dict'): if hasattr(data, 'to_dict'):
return kombu.utils.json.dumps(dict( native = dict(
type='dataclass_json', type='dataclass_json',
module=data.__class__.__module__, module=data.__class__.__module__,
_class=data.__class__.__qualname__, _class=data.__class__.__qualname__,
data=data.to_dict(), data=data.to_dict(),
)) )
else: else:
return kombu.utils.json.dumps(data) if isinstance(data, (list, tuple)):
native = [
cls._dumps(o, need_native=True,)
for o in data
]
elif isinstance(data, dict):
native = {
k : cls._dumps(v, need_native=True,)
for k, v in data.items()
}
else:
native = data
if not need_native:
return kombu.utils.json.dumps(native)
else:
return native
def kombu_register_json_dataclass(): def kombu_register_json_dataclass():
import kombu.serialization import kombu.serialization
@ -47,5 +88,5 @@ def kombu_register_json_dataclass():
Task._dumps, Task._dumps,
Task._loads, Task._loads,
content_type='application/json', content_type='application/json',
content_encoding='text', content_encoding='utf-8',
) )