def test_track_functions(self):
assert _is_function(f1), "function got decorated unexpectedly"
assert not _is_task(f1), "function got decorated unexpectedly"
track_functions(f1, f2, f3)
assert callable(f1), "function is not function anymore"
assert _is_task(f1), "local function wasn't decorated"
assert _is_task(
module_to_track.f1), "function in module wasn't decorated"
assert _is_task(f2)
assert _is_task(f3)
def test_track_modules(self):
assert _is_function(f4), "function got decorated unexpectedly"
assert not _is_task(f4), "function got decorated unexpectedly"
track_modules(module_to_track)
assert callable(f4), "function is not function anymore"
assert _is_task(f4), "local function wasn't decorated"
assert callable(module_to_track.f4), "function is not function anymore"
assert _is_task(
module_to_track.f4), "function in module wasn't decorated"
assert _is_task(module_to_track.f5)
assert _is_function(module_to_track.task)
assert not _is_task(module_to_track.task)
def _get_task_cls(self, task_name):
from dbnd._core.utils.basics.load_python_module import load_python_module
task_cls = self._get_registered_task_cls(task_name)
if task_cls:
return task_cls
# we are going to check if we have override/definition in config
config_task_type = config.get(task_name, "_type", None)
if config_task_type:
_validate_no_recursion_in_config(task_name, config_task_type,
"_type")
try:
return self._get_task_cls(config_task_type)
except Exception:
logger.warning(
"Failed to load type required by [%s] using _type=%s",
task_name,
config_task_type,
)
raise
config_task_type = config.get(task_name, "_from", None)
if config_task_type:
_validate_no_recursion_in_config(task_name, config_task_type,
"_from")
return self._get_task_cls(config_task_type)
if "." in task_name:
parts = task_name.split(".")
possible_root_task = parts.pop()
possible_module = ".".join(parts)
# Try to load module and check again for existance
load_python_module(possible_module, "task name '%s'" % task_name)
task_cls = self._get_registered_task_cls(task_name)
if task_cls:
return task_cls
# Check if task exists but user forgot to decorate method with @task
task_module = sys.modules.get(possible_module)
if task_module and hasattr(task_module, possible_root_task):
user_func = getattr(task_module, possible_root_task)
from dbnd._core.task_build.task_metaclass import TaskMetaclass
# weird case of we didn't manage to find a class after it was loaded via _get_registered_task_cls
# can happen if the function name is different from "task_family"
if isinstance(user_func, TaskMetaclass):
return user_func
if callable(user_func):
from dbnd._core.tracking.python_tracking import _is_task
if _is_task(user_func):
return user_func.task_cls
# Non-decorated function was found - decorate and return it
from dbnd._core.task_build import dbnd_decorator
decorated_task = dbnd_decorator.task(user_func)
setattr(task_module, possible_root_task, decorated_task)
logger.warning(
"Found non-decorated task: %s. "
"Please decorate this task with the proper symbol @pipeline / @task.\n"
"Auto-decorating and treating it as @task ...",
task_name,
)
return decorated_task.task
if is_airflow_enabled():
from dbnd_airflow.dbnd_task_executor.airflow_operator_as_dbnd import (
AirflowDagAsDbndTask, )
dag = self._get_aiflow_dag(task_name)
if dag:
return AirflowDagAsDbndTask
return None