gokart.build(reset_register=False) can use PandasTypeConfig

gokart/build.py

@@ -4,6 +4,7 @@
 
 import luigi
 
+import gokart
 from gokart.task import TaskOnKart
 
 
@@ -38,8 +39,11 @@ def _get_output(task: TaskOnKart) -> Any:
 
 
 def _reset_register(keep={'gokart', 'luigi'}):
-    luigi.task_register.Register._reg = [x for x in luigi.task_register.Register._reg
-                                         if x.__module__.split('.')[0] in keep]  # avoid TaskClassAmbigiousException
+    """reset luigi.task_register.Register._reg everytime gokart.build called to avoid TaskClassAmbigiousException"""
+    luigi.task_register.Register._reg = [
+        x for x in luigi.task_register.Register._reg if ((x.__module__.split('.')[0] in keep)  # keep luigi and gokart
+                                                         or (issubclass(x, gokart.PandasTypeConfig)))  # PandasTypeConfig should be kept
+    ]
 
 
 def build(task: TaskOnKart, return_value: bool = True, reset_register: bool = True, log_level: int = logging.ERROR, **env_params) -> Optional[Any]:

pyproject.toml

@@ -48,7 +48,8 @@ types-redis = "*"
 [tool.flake8]
 # B006: Do not use mutable data structures for argument defaults. They are created during function definition time. All calls to the function reuse this one instance of that data structure, persisting changes between them.
 # B008 Do not perform function calls in argument defaults.  The call is performed only once at function definition time. All calls to your function will reuse the result of that definition-time function call.  If this is intended, assign the function call to a module-level variable and use that variable as a default value.
-ignore = "B006,B008"
+# W503 line break before binary operator. We use W504(line break after binary operator) rather than W503
+ignore = "B006,B008,W503"
 max-line-length = 160
 exclude = "venv/*,tox/*"
 

test/test_pandas_type_check_framework.py

@@ -1,3 +1,4 @@
+import logging
 import unittest
 from logging import getLogger
 from typing import Any, Dict
@@ -8,6 +9,7 @@
 from luigi.mock import MockFileSystem, MockTarget
 
 import gokart
+from gokart.build import GokartBuildError
 from gokart.pandas_type_config import PandasTypeConfig
 
 logger = getLogger(__name__)
@@ -63,28 +65,29 @@ def setUp(self) -> None:
         luigi.setup_logging.DaemonLogging._configured = False
         luigi.setup_logging.InterfaceLogging._configured = False
         MockFileSystem().clear()
+        # same way as luigi https://github.com/spotify/luigi/blob/fe7ecf4acf7cf4c084bd0f32162c8e0721567630/test/helpers.py#L175
+        self._stashed_reg = luigi.task_register.Register._get_reg()
 
     def tearDown(self) -> None:
         luigi.setup_logging.DaemonLogging._configured = False
         luigi.setup_logging.InterfaceLogging._configured = False
+        luigi.task_register.Register._set_reg(self._stashed_reg)
 
     @patch('sys.argv', new=['main', 'test_pandas_type_check_framework._DummyFailTask', '--log-level=CRITICAL', '--local-scheduler', '--no-lock'])
     @patch('luigi.LocalTarget', new=lambda path, **kwargs: MockTarget(path, **kwargs))
-    def test_fail(self):
+    def test_fail_with_gokart_run(self):
         with self.assertRaises(SystemExit) as exit_code:
             gokart.run()
         self.assertNotEqual(exit_code.exception.code, 0)  # raise Error
 
-    @patch('sys.argv', new=['main', 'test_pandas_type_check_framework._DummyFailWithNoneTask', '--log-level=CRITICAL', '--local-scheduler', '--no-lock'])
-    @patch('luigi.LocalTarget', new=lambda path, **kwargs: MockTarget(path, **kwargs))
+    def test_fail(self):
+        with self.assertRaises(GokartBuildError):
+            gokart.build(_DummyFailTask(), log_level=logging.CRITICAL)
+
     def test_fail_with_None(self):
-        with self.assertRaises(SystemExit) as exit_code:
-            gokart.run()
-        self.assertNotEqual(exit_code.exception.code, 0)  # raise Error
+        with self.assertRaises(GokartBuildError):
+            gokart.build(_DummyFailWithNoneTask(), log_level=logging.CRITICAL)
 
-    @patch('sys.argv', new=['main', 'test_pandas_type_check_framework._DummySuccessTask', '--log-level=CRITICAL', '--local-scheduler', '--no-lock'])
-    @patch('luigi.LocalTarget', new=lambda path, **kwargs: MockTarget(path, **kwargs))
     def test_success(self):
-        with self.assertRaises(SystemExit) as exit_code:
-            gokart.run()
-        self.assertEqual(exit_code.exception.code, 0)
+        gokart.build(_DummySuccessTask())
+        # no error