Fpin/refactor action file loading

karadoc/common/commands/command.py

@@ -48,7 +48,7 @@ def add_arguments(parser: ArgumentParser) -> None:
     @abstractmethod
     def do_command(args: Namespace) -> Optional[ReturnCode]:
         """Entrypoint of the command.
-        This method can return nothing if it's only way of failing is raising an Exception.
+        This method can return nothing if its only way of failing is raising an Exception.
         Commands that can fail without raising exceptions (e.g. a validation command) should return a ReturnCode.
         """
         pass

karadoc/common/job_core/has_before_after.py

@@ -1,39 +1,41 @@
+from karadoc.common.job_core.package import OptionalMethod
+
+
 class HasBeforeAfter:
-    def __init__(self) -> None:
-        # Attributes that the user is not supposed to change
-        self.__before_each = None
-        self.__after_each = None
-        self.__before_all = None
-        self.__after_all = None
+    def __init__(self):
+        def default() -> None:
+            # We do nothing by default
+            pass
+
+        self.before_each = OptionalMethod(default, "before_each")
+        self.after_each = OptionalMethod(default, "after_each")
+        self.before_all = OptionalMethod(default, "before_all")
+        self.after_all = OptionalMethod(default, "after_all")
 
     def before_each(self) -> None:
         """This method will be called before executing each action method defined in the job's definition file.
 
         This is similar to unittest's `setup()` method.
         """
-        if self.__before_each is not None:
-            return self.__before_each()
+        pass
 
     def after_each(self) -> None:
         """This method will be called after executing each action method defined in the job's definition file.
 
         This is similar to unittest's `teardown()` method.
         """
-        if self.__after_each is not None:
-            return self.__after_each()
+        pass
 
     def before_all(self) -> None:
         """This method will be called once before executing all the action methods defined in the job's definition file.
 
         This is similar to unittest's `setupClass()` method.
         """
-        if self.__before_all is not None:
-            return self.__before_all()
+        pass
 
     def after_all(self) -> None:
         """This method will be called once after executing all the action methods defined in the job's definition file.
 
         This is similar to unittest's `teardownClass()` method.
         """
-        if self.__after_all is not None:
-            return self.__after_all()
+        pass

karadoc/common/job_core/job_base.py

@@ -1,11 +1,8 @@
-from typing import Optional
-
 from karadoc.common.utils.assert_utils import assert_true
 
 
 class JobBase:
     _action_file_name_conf_key: str
-    _run_method_name: Optional[str] = None
     output: str
 
     @classmethod

karadoc/common/job_core/load.py

@@ -4,36 +4,17 @@
 from os.path import isfile, splitext
 from pathlib import Path
 from types import FunctionType, ModuleType
-from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, cast
+from typing import Dict, Optional, Type, TypeVar, Union, cast
 
 from karadoc.common import conf
 from karadoc.common.exceptions import ActionFileLoadingError, ForbiddenActionError
-from karadoc.common.job_core.has_before_after import HasBeforeAfter
 from karadoc.common.job_core.has_vars import HasVars
 from karadoc.common.job_core.job_base import JobBase
+from karadoc.common.job_core.package import ActionFileMethod
 from karadoc.common.model import file_index
 from karadoc.common.table_utils import parse_table_name
 from karadoc.common.utils.assert_utils import assert_true
-from karadoc.spark.job_core.has_external_inputs import (
-    HasExternalInputs,
-    _read_external_input_signature_check,
-    _read_external_inputs_signature_check,
-)
-from karadoc.spark.job_core.has_external_outputs import (
-    HasExternalOutputs,
-    _write_external_output_signature_check,
-    _write_external_outputs_signature_check,
-)
 from karadoc.spark.job_core.has_spark import HasSpark
-from karadoc.spark.job_core.has_stream_external_inputs import (
-    HasStreamExternalInputs,
-    _read_stream_external_input_signature_check,
-    _read_stream_external_inputs_signature_check,
-)
-from karadoc.spark.job_core.has_stream_external_output import (
-    HasStreamExternalOutput,
-    _write_stream_external_output_signature_check,
-)
 from karadoc.spark.quality.checks import Alert, Metric
 
 Job = TypeVar("Job", JobBase, JobBase)
@@ -58,15 +39,7 @@ def load_non_runnable_action_file(full_table_name: str, job_type: Type[Job]) ->
     passed_vars = None
 
     job = __load_action_file(job_type, full_table_name, passed_vars)
-    if job._run_method_name is not None:
-
-        def empty_run() -> None:
-            raise ForbiddenActionError(
-                f"The {job._run_method_name} method of a job returned by the `load_non_runnable_action_file` "
-                "method cannot be called. use load_runnable_action_file instead."
-            )
-
-        job.__setattr__(job._run_method_name, empty_run)
+    __unset_all_action_methods(job)
     return job
 
 
@@ -175,46 +148,35 @@ def __load_module_file(module_name: str, module_path: Union[str, Path]) -> Modul
     return mod
 
 
-def __load_optional_method(
-    mod: ModuleType, tpe: type, method_name: str, signature_method: Optional[Callable[..., Any]] = None
-) -> None:
-    if hasattr(mod, method_name):
-        private_method_name = "_" + tpe.__name__ + "__" + method_name
-        if signature_method is not None:
-            signature_method = cast(FunctionType, signature_method)
-            check_method_signatures(method_name, getattr(mod, method_name), signature_method)
-        setattr(mod.job, private_method_name, getattr(mod, method_name))
-
-
-def _set_spark_batch_job(mod: ModuleType) -> None:
-    from karadoc.spark.batch.spark_batch_job import SparkBatchJob
-
-    if isinstance(mod.job, SparkBatchJob):
-        mod.job.run = mod.run
+def __load_action_methods(mod: ModuleType) -> None:
+    action_methods = [method for name, method in inspect.getmembers(mod.job) if isinstance(method, ActionFileMethod)]
+    for method in action_methods:
+        if hasattr(mod, method.name):
+            method_defined_in_file = getattr(mod, method.name)
+            method.set_method(method_defined_in_file)
 
 
-def _set_spark_stream_job(mod: ModuleType) -> None:
-    from karadoc.spark.stream.spark_stream_job import SparkStreamJob
+def __unset_action_method(job: Job, method: ActionFileMethod):
+    def empty_run() -> None:
+        raise ForbiddenActionError(
+            f"The {method.name} method of a job returned by the `load_non_runnable_action_file` "
+            "method cannot be called. Use load_runnable_action_file instead."
+        )
 
-    if isinstance(mod.job, SparkStreamJob):
-        mod.job.stream = mod.stream
+    job.__setattr__(method.name, empty_run)
 
 
-def _set_analyze_job(mod: ModuleType) -> None:
-    from karadoc.spark.analyze.analyze_job import AnalyzeJob
+def __unset_all_action_methods(job: Job) -> None:
+    action_methods = [method for name, method in inspect.getmembers(job) if isinstance(method, ActionFileMethod)]
 
-    if isinstance(mod.job, AnalyzeJob):
-        mod.job.analyze = mod.analyze
+    for method in action_methods:
+        __unset_action_method(job, method)
 
 
 def _set_quality_check_job(mod: ModuleType) -> None:
     from karadoc.spark.quality.quality_check_job import QualityCheckJob
 
     if isinstance(mod.job, QualityCheckJob):
-        __load_optional_method(mod, HasBeforeAfter, "before_all")
-        __load_optional_method(mod, HasBeforeAfter, "after_all")
-        __load_optional_method(mod, HasBeforeAfter, "before_each")
-        __load_optional_method(mod, HasBeforeAfter, "after_each")
         found_alerts = [obj for name, obj in inspect.getmembers(mod) if isinstance(obj, Alert)]
         for alert in found_alerts:
             mod.job.add_alert(alert)
@@ -223,48 +185,9 @@ def _set_quality_check_job(mod: ModuleType) -> None:
             mod.job.add_metric(metric)
 
 
-def _set_has_external_inputs(mod: ModuleType) -> None:
-    if isinstance(mod.job, HasExternalInputs):
-        __load_optional_method(mod, HasExternalInputs, "read_external_input", _read_external_input_signature_check())
-        __load_optional_method(mod, HasExternalInputs, "read_external_inputs", _read_external_inputs_signature_check())
-
-
-def _set_has_stream_external_inputs(mod: ModuleType) -> None:
-    if isinstance(mod.job, HasStreamExternalInputs):
-        __load_optional_method(
-            mod, HasStreamExternalInputs, "read_external_input", _read_stream_external_input_signature_check()
-        )
-        __load_optional_method(
-            mod, HasStreamExternalInputs, "read_external_inputs", _read_stream_external_inputs_signature_check()
-        )
-
-
-def _set_has_external_outputs(mod: ModuleType) -> None:
-    if isinstance(mod.job, HasExternalOutputs):
-        __load_optional_method(
-            mod, HasExternalOutputs, "write_external_output", _write_external_output_signature_check()
-        )
-        __load_optional_method(
-            mod, HasExternalOutputs, "write_external_outputs", _write_external_outputs_signature_check()
-        )
-
-
-def _set_has_stream_external_output(mod: ModuleType) -> None:
-    if isinstance(mod.job, HasStreamExternalOutput):
-        __load_optional_method(
-            mod, HasStreamExternalOutput, "write_external_output", _write_stream_external_output_signature_check()
-        )
-
-
 def __set_job_from_module(mod: ModuleType) -> JobBase:
-    _set_spark_batch_job(mod)
-    _set_spark_stream_job(mod)
-    _set_analyze_job(mod)
+    __load_action_methods(mod)
     _set_quality_check_job(mod)
-    _set_has_external_inputs(mod)
-    _set_has_stream_external_inputs(mod)
-    _set_has_external_outputs(mod)
-    _set_has_stream_external_output(mod)
     return mod.job
 
 
@@ -281,8 +204,8 @@ def __load_file(
     file_path: str,
     file_type: str,
     passed_vars: Optional[Dict[str, str]],
-    job_type: type,
-) -> JobBase:
+    job_type: Type[Job],
+) -> Job:
     full_table_name = schema_name + "." + table_name
     if file_path is None:
         raise ActionFileLoadingError(f"Could not find a {file_type} for table {full_table_name}")

karadoc/common/job_core/package.py

@@ -0,0 +1,80 @@
+import inspect
+from abc import ABC, abstractmethod
+from types import FunctionType
+from typing import Callable, Optional
+
+
+def method_signature_str(method_name: str, full_arg_spec: inspect.FullArgSpec) -> str:
+    annotations = full_arg_spec.annotations
+
+    def arg_str(arg_name: str) -> str:
+        if arg_name in annotations:
+            return f"{arg_name}: {annotations[arg_name]}"
+        else:
+            return arg_name
+
+    args = [arg_str(arg_name) for arg_name in full_arg_spec.args]
+    return_str = f" -> {annotations['return']}" if "return" in annotations else ""
+    return f"def {method_name}({', '.join(args)}){return_str}:"
+
+
+def check_method_signatures(method_name: str, actual: Callable, expected: Callable) -> None:
+    """Ensures that the signature of a given method matches the expected signature taken from another method.
+
+    :param method_name: name of the method to check
+    :param actual: function to check
+    :param expected: control function with the expected signature
+    :return:
+    """
+    if not isinstance(actual, FunctionType):
+        raise TypeError("%s is not a function" % method_name)
+    actual_args = inspect.getfullargspec(actual)
+    expected_args = inspect.getfullargspec(expected)
+    if actual_args != expected_args:
+        raise TypeError(
+            f"The method {method_name} should have the following signature\n"
+            + method_signature_str(method_name, expected_args)
+        )
+
+
+class ActionFileMethod(ABC):
+    name: str
+
+    @abstractmethod
+    def __call__(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def set_method(self, func: Callable) -> None:
+        """Description of the command which will be displayed in the help"""
+        pass
+
+
+class OptionalMethod(ActionFileMethod):
+    def __init__(self, default_func: Callable, name: Optional[str] = None):
+        self.func = default_func
+        if name is None:
+            name = default_func.__name__
+        self.name = name
+
+    def __call__(self, *args, **kwargs):
+        return self.func(*args, **kwargs)
+
+    def set_method(self, func: Callable) -> None:
+        check_method_signatures(self.name, func, self.func)
+        self.func = func
+
+
+class RequiredMethod(ActionFileMethod):
+    def __init__(self, name: str, signature_func: Optional[Callable] = None):
+        self.name = name
+        self.signature_func = signature_func
+        self.func = None
+
+    def __call__(self, *args, **kwargs) -> None:
+        return self.func(*args, **kwargs)
+
+    def set_method(self, func: Callable) -> None:
+        if self.signature_func is not None:
+            check_method_signatures(self.name, func, self.signature_func)
+        self.func = func

karadoc/spark/analyze/analyze_job.py

@@ -1,11 +1,11 @@
 from karadoc.common.job_core.has_vars import HasVars
+from karadoc.common.job_core.package import RequiredMethod
 from karadoc.spark.job_core.has_batch_inputs import HasBatchInputs
 from karadoc.spark.job_core.has_spark import HasSpark
 
 
 class AnalyzeJob(HasBatchInputs, HasVars, HasSpark):
     _action_file_name_conf_key = "spark.analyze_timeline"
-    _run_method_name = "analyze"
 
     def __init__(self) -> None:
         HasSpark.__init__(self)
@@ -16,4 +16,5 @@ def __init__(self) -> None:
         self.reference_time_col = "application_date"
         self.cohorts = ["cohort"]
         self.nb_buckets = 5
-        self.analyze = None
+
+        self.analyze = RequiredMethod("analyze")

karadoc/spark/batch/spark_batch_job.py

@@ -1,6 +1,7 @@
 from karadoc.common.job_core.has_disable import HasDisable
 from karadoc.common.job_core.has_keys import HasKeys
 from karadoc.common.job_core.has_vars import HasVars
+from karadoc.common.job_core.package import RequiredMethod
 from karadoc.spark.job_core.has_batch_inputs import HasBatchInputs
 from karadoc.spark.job_core.has_batch_output import HasBatchOutput
 from karadoc.spark.job_core.has_external_inputs import HasExternalInputs
@@ -19,7 +20,6 @@ class SparkBatchJob(
     HasSpark,
 ):
     _action_file_name_conf_key = "spark.batch"
-    _run_method_name = "run"
 
     def __init__(self) -> None:
         HasSpark.__init__(self)
@@ -31,4 +31,4 @@ def __init__(self) -> None:
         HasKeys.__init__(self)
         HasDisable.__init__(self)
 
-        self.run = None
+        self.run = RequiredMethod("run")