From 26b46be6de226f5db3d4975816908053446aa23b Mon Sep 17 00:00:00 2001
From: PuQing <me@puqing.work>
Date: Sun, 28 Jan 2024 16:33:04 +0800
Subject: [PATCH] Update imports and auto register tasks in worker.py

---
 backend/app/app/source/ICLR.py     | 2 +-
 backend/app/app/source/ICML.py     | 2 +-
 backend/app/app/source/__init__.py | 3 +--
 backend/app/app/worker.py          | 9 ++++++---
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/backend/app/app/source/ICLR.py b/backend/app/app/source/ICLR.py
index e70dfe7..eb1d6f9 100644
--- a/backend/app/app/source/ICLR.py
+++ b/backend/app/app/source/ICLR.py
@@ -1,4 +1,4 @@
-from app.source import NIPS
+from app.source.NIPS import NIPS
 
 
 class ICLR(NIPS):
diff --git a/backend/app/app/source/ICML.py b/backend/app/app/source/ICML.py
index 967a005..84a907a 100644
--- a/backend/app/app/source/ICML.py
+++ b/backend/app/app/source/ICML.py
@@ -1,4 +1,4 @@
-from app.source import NIPS
+from app.source.NIPS import NIPS
 
 
 class ICML(NIPS):
diff --git a/backend/app/app/source/__init__.py b/backend/app/app/source/__init__.py
index 901d986..a934920 100644
--- a/backend/app/app/source/__init__.py
+++ b/backend/app/app/source/__init__.py
@@ -1,7 +1,6 @@
 from app.source.Arxiv import Arxiv
-from app.source.base import PaperRequestsTask
 from app.source.ICLR import ICLR
 from app.source.ICML import ICML
 from app.source.NIPS import NIPS
 
-__all__ = ["PaperRequestsTask", "NIPS", "ICLR", "ICML", "Arxiv"]
+__all__ = ["NIPS", "ICLR", "ICML", "Arxiv"]
diff --git a/backend/app/app/worker.py b/backend/app/app/worker.py
index c0f042b..d653c5f 100644
--- a/backend/app/app/worker.py
+++ b/backend/app/app/worker.py
@@ -1,13 +1,16 @@
+import inspect
 from typing import Union
 
 from celery import Task
 from sqlmodel import Session
 
+from app import source
 from app.core.celery_app import celery_app
-from app.source import NIPS, Arxiv
 
-celery_app.register_task(NIPS())
-celery_app.register_task(Arxiv())
+members = inspect.getmembers(source, inspect.isclass)
+for name, _class in members:
+    # auto register task
+    celery_app.register_task(_class)
 
 
 def batch(iterable: Union[set[str], list[str]], n: int = 1):