diff --git a/gokart/build.py b/gokart/build.py index 7662cfa5..68a3a7a9 100644 --- a/gokart/build.py +++ b/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]: diff --git a/pyproject.toml b/pyproject.toml index a05a29a1..7010e26d 100644 --- a/pyproject.toml +++ b/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/*" diff --git a/test/test_pandas_type_check_framework.py b/test/test_pandas_type_check_framework.py index c36fb99d..4ec92b2c 100644 --- a/test/test_pandas_type_check_framework.py +++ b/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