From ff16d48f869ce65f9ff2ed8c3bfcebb28d936b2c Mon Sep 17 00:00:00 2001
From: Siarhei Siniak <siarheisiniak@gmail.com>
Date: Sat, 6 Jul 2024 19:54:47 +0300
Subject: [PATCH] [~] Refactor

---
 python/tasks/tiktok/tasks.py |  6 +++---
 python/tasks/tiktok/utils.py | 12 ++++++++++++
 2 files changed, 15 insertions(+), 3 deletions(-)
 create mode 100644 python/tasks/tiktok/utils.py

diff --git a/python/tasks/tiktok/tasks.py b/python/tasks/tiktok/tasks.py
index b803921..49f2b49 100644
--- a/python/tasks/tiktok/tasks.py
+++ b/python/tasks/tiktok/tasks.py
@@ -16,6 +16,7 @@ from typing import (
 )
 import celery
 from .config import tiktok_config
+from .utils import Task
 
 #logging.getLogger().setLevel(logging.INFO)
 
@@ -307,8 +308,8 @@ def tiktok_videos_process(meta: Iterable[dict[str, Any]]) -> dict[str, Any]:
     return stats
 
 class audio_get_t:
-    @dataclasses.dataclass
     @dataclasses_json.dataclass_json
+    @dataclasses.dataclass
     class res_t:
         file: str
         file_mp3: str
@@ -316,7 +317,7 @@ class audio_get_t:
         path_mp3: str
         url: str
 
-@celery.shared_task()
+@celery.shared_task(base=Task)
 def audio_get() -> audio_get_t.res_t:
     c = tiktok_config()
     url = 'https://www.youtube.com/watch?v=dQw4w9WgXcQ'
@@ -348,4 +349,3 @@ def audio_get() -> audio_get_t.res_t:
         path_mp3=path_mp3,
         url=url,
     )
-
diff --git a/python/tasks/tiktok/utils.py b/python/tasks/tiktok/utils.py
new file mode 100644
index 0000000..9c224fe
--- /dev/null
+++ b/python/tasks/tiktok/utils.py
@@ -0,0 +1,12 @@
+import celery.app.task
+
+class Task(celery.app.task.Task):
+    def __call__(self, *args, **kwargs):
+        res = super().__call__(*args, **kwargs)
+        return self._to_native(res)
+
+    def _to_native(self, data):
+        if hasattr(data, 'to_dict'):
+            return data.to_dict()
+        else:
+            return data