fix import style for redis.lock and luigi.task_register

gokart/build.py

@@ -10,7 +10,7 @@
 
 import backoff
 import luigi
-from luigi import LuigiStatusCode, rpc, scheduler
+from luigi import LuigiStatusCode, rpc, scheduler, task_register
 
 import gokart
 import gokart.tree.task_info
@@ -95,10 +95,10 @@ def _get_output(task: TaskOnKart[T]) -> T:
 
 
 def _reset_register(keep={'gokart', 'luigi'}):
-    """reset luigi.task_register.Register._reg everytime gokart.build called to avoid TaskClassAmbigiousException"""
-    luigi.task_register.Register._reg = [
+    """reset task_register.Register._reg everytime gokart.build called to avoid TaskClassAmbigiousException"""
+    task_register.Register._reg = [
         x
-        for x in luigi.task_register.Register._reg
+        for x in task_register.Register._reg
         if (
             (x.__module__.split('.')[0] in keep)  # keep luigi and gokart
             or (issubclass(x, gokart.PandasTypeConfig))

gokart/conflict_prevention_lock/task_lock.py

@@ -7,6 +7,7 @@
 
 import redis
 from apscheduler.schedulers.background import BackgroundScheduler
+from redis.lock import Lock
 
 logger = getLogger(__name__)
 
@@ -47,20 +48,20 @@ def get_redis_client(self):
         return self._redis_client
 
 
-def _extend_lock(task_lock: redis.lock.Lock, redis_timeout: int) -> None:
+def _extend_lock(task_lock: Lock, redis_timeout: int) -> None:
     task_lock.extend(additional_time=redis_timeout, replace_ttl=True)
 
 
-def set_task_lock(task_lock_params: TaskLockParams) -> redis.lock.Lock:
+def set_task_lock(task_lock_params: TaskLockParams) -> Lock:
     redis_client = RedisClient(host=task_lock_params.redis_host, port=task_lock_params.redis_port).get_redis_client()
     blocking = not task_lock_params.raise_task_lock_exception_on_collision
-    task_lock = redis.lock.Lock(redis=redis_client, name=task_lock_params.redis_key, timeout=task_lock_params.redis_timeout, thread_local=False)
+    task_lock = Lock(redis=redis_client, name=task_lock_params.redis_key, timeout=task_lock_params.redis_timeout, thread_local=False)
     if not task_lock.acquire(blocking=blocking):
         raise TaskLockException('Lock already taken by other task.')
     return task_lock
 
 
-def set_lock_scheduler(task_lock: redis.lock.Lock, task_lock_params: TaskLockParams) -> BackgroundScheduler:
+def set_lock_scheduler(task_lock: Lock, task_lock_params: TaskLockParams) -> BackgroundScheduler:
     scheduler = BackgroundScheduler()
     extend_lock = functools.partial(_extend_lock, task_lock=task_lock, redis_timeout=task_lock_params.redis_timeout or 0)
     scheduler.add_job(