Skip to content

Commit

Permalink
gokart.build(reset_register=False) can use PandasTypeConfig (#335)
Browse files Browse the repository at this point in the history
* ✅ test PandasTypeCheck and build

* 🎨 do not reset PandasTypeConfig when reset_register

* 🎨 better register reset in tests 'https://github.com/spotify/luigi/blob/fe7ecf4acf7cf4c084bd0f32162c8e0721567630/test/helpers.py#L175'

* 👕 isort and plake8

* 👕 mute ERRORS on test

---------

Co-authored-by: Keisuke Ogaki <[email protected]>
  • Loading branch information
Hi-king and Keisuke Ogaki authored Nov 21, 2023
1 parent 919b4c7 commit 5f37e57
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 14 deletions.
8 changes: 6 additions & 2 deletions gokart/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import luigi

import gokart
from gokart.task import TaskOnKart


Expand Down Expand Up @@ -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]:
Expand Down
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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/*"

Expand Down
25 changes: 14 additions & 11 deletions test/test_pandas_type_check_framework.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import logging
import unittest
from logging import getLogger
from typing import Any, Dict
Expand All @@ -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__)
Expand Down Expand Up @@ -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

0 comments on commit 5f37e57

Please sign in to comment.