fractal-task-tools 0.2.0__py3-none-any.whl

fractal_task_tools/_descriptions.py

@@ -0,0 +1,252 @@
+import ast
+import logging
+import os
+from importlib import import_module
+from pathlib import Path
+from typing import Optional
+
+from docstring_parser import parse as docparse
+
+
+def _sanitize_description(string: str) -> str:
+    """
+    Sanitize a description string.
+
+    This is a provisional helper function that replaces newlines with spaces
+    and reduces multiple contiguous whitespace characters to a single one.
+    Future iterations of the docstrings format/parsing may render this function
+    not-needed or obsolete.
+
+    Args:
+        string: TBD
+    """
+    # Replace newline with space
+    new_string = string.replace("\n", " ")
+    # Replace N-whitespace characters with a single one
+    while "  " in new_string:
+        new_string = new_string.replace("  ", " ")
+    return new_string
+
+
+def _get_function_docstring(
+    *,
+    package_name: Optional[str],
+    module_path: str,
+    function_name: str,
+    verbose: bool = False,
+) -> str:
+    """
+    Extract docstring from a function.
+
+
+    Args:
+        package_name: Example `fractal_tasks_core`.
+        module_path:
+            This must be an absolute path like `/some/module.py` (if
+            `package_name` is `None`) or a relative path like `something.py`
+            (if `package_name` is not `None`).
+        function_name: Example `create_ome_zarr`.
+    """
+
+    if not module_path.endswith(".py"):
+        raise ValueError(f"Module {module_path} must end with '.py'")
+
+    # Get the function ast.FunctionDef object
+    if package_name is not None:
+        if os.path.isabs(module_path):
+            raise ValueError(
+                "Error in _get_function_docstring: `package_name` is not "
+                "None but `module_path` is absolute."
+            )
+        package_path = Path(import_module(package_name).__file__).parent
+        module_path = package_path / module_path
+    else:
+        if not os.path.isabs(module_path):
+            raise ValueError(
+                "Error in _get_function_docstring: `package_name` is None "
+                "but `module_path` is not absolute."
+            )
+        module_path = Path(module_path)
+
+    if verbose:
+        logging.info(f"[_get_function_docstring] {function_name=}")
+        logging.info(f"[_get_function_docstring] {module_path=}")
+
+    tree = ast.parse(module_path.read_text())
+    _function = next(
+        f
+        for f in ast.walk(tree)
+        if (isinstance(f, ast.FunctionDef) and f.name == function_name)
+    )
+
+    # Extract docstring from ast.FunctionDef
+    return ast.get_docstring(_function)
+
+
+def _get_function_args_descriptions(
+    *,
+    package_name: Optional[str],
+    module_path: str,
+    function_name: str,
+    verbose: bool = False,
+) -> dict[str, str]:
+    """
+    Extract argument descriptions from a function.
+
+    Args:
+        package_name: Example `fractal_tasks_core`.
+        module_path:
+            This must be an absolute path like `/some/module.py` (if
+            `package_name` is `None`) or a relative path like `something.py`
+            (if `package_name` is not `None`).
+        function_name: Example `create_ome_zarr`.
+    """
+
+    # Extract docstring from ast.FunctionDef
+    docstring = _get_function_docstring(
+        package_name=package_name,
+        module_path=module_path,
+        function_name=function_name,
+        verbose=verbose,
+    )
+    if verbose:
+        logging.info(f"[_get_function_args_descriptions] {docstring}")
+
+    # Parse docstring (via docstring_parser) and prepare output
+    parsed_docstring = docparse(docstring)
+    descriptions = {
+        param.arg_name: _sanitize_description(param.description)
+        for param in parsed_docstring.params
+    }
+    logging.info(f"[_get_function_args_descriptions] END ({function_name=})")
+    return descriptions
+
+
+def _get_class_attrs_descriptions_from_file(
+    *,
+    module_path: Path,
+    class_name: str,
+) -> dict[str, str]:
+    """
+    Extract class-attribute descriptions from a Python script
+
+    Args:
+        module_path: Example `/something/my_class.py`.
+        class_name: Example `OmeroChannel`.
+    """
+    tree = ast.parse(module_path.read_text())
+    try:
+        _class = next(
+            c
+            for c in ast.walk(tree)
+            if (isinstance(c, ast.ClassDef) and c.name == class_name)
+        )
+    except StopIteration:
+        raise RuntimeError(f"Cannot find {class_name=} in {module_path}.")
+    docstring = ast.get_docstring(_class)
+    parsed_docstring = docparse(docstring)
+    descriptions = {
+        x.arg_name: _sanitize_description(x.description)
+        if x.description
+        else "Missing description"
+        for x in parsed_docstring.params
+    }
+    return descriptions
+
+
+def _get_class_attrs_descriptions(
+    package_name: str, module_relative_path: str, class_name: str
+) -> dict[str, str]:
+    """
+    Extract class-attribute descriptions from an imported module
+
+    Args:
+        package_name: Example `fractal_tasks_core`.
+        module_relative_path: Example `lib_channels.py`.
+        class_name: Example `OmeroChannel`.
+    """
+
+    if not module_relative_path.endswith(".py"):
+        raise ValueError(f"Module {module_relative_path} must end with '.py'")
+
+    # Get the class ast.ClassDef object
+    package_path = Path(import_module(package_name).__file__).parent
+    module_path = package_path / module_relative_path
+    descriptions = _get_class_attrs_descriptions_from_file(
+        module_path=module_path,
+        class_name=class_name,
+    )
+    logging.info(f"[_get_class_attrs_descriptions] END ({class_name=})")
+    return descriptions
+
+
+def _insert_function_args_descriptions(
+    *, schema: dict, descriptions: dict, verbose: bool = False
+):
+    """
+    Merge the descriptions obtained via `_get_args_descriptions` into the
+    properties of an existing JSON Schema.
+
+    Args:
+        schema: TBD
+        descriptions: TBD
+    """
+    new_schema = schema.copy()
+    new_properties = schema["properties"].copy()
+    for key, value in schema["properties"].items():
+        if "description" in value:
+            raise ValueError("Property already has description")
+        else:
+            if key in descriptions:
+                value["description"] = descriptions[key]
+            else:
+                value["description"] = "Missing description"
+            new_properties[key] = value
+            if verbose:
+                logging.info(
+                    "[_insert_function_args_descriptions] "
+                    f"Add {key=}, {value=}"
+                )
+    new_schema["properties"] = new_properties
+    logging.info("[_insert_function_args_descriptions] END")
+    return new_schema
+
+
+def _insert_class_attrs_descriptions(
+    *,
+    schema: dict,
+    class_name: str,
+    descriptions: dict,
+    definition_key: str,
+):
+    """
+    Merge the descriptions obtained via `_get_attributes_models_descriptions`
+    into the `class_name` definition, within an existing JSON Schema
+
+    Args:
+        schema: TBD
+        class_name: TBD
+        descriptions: TBD
+        definition_key: Either `"definitions"` (for Pydantic V1) or
+            `"$defs"` (for Pydantic V2)
+    """
+    new_schema = schema.copy()
+    if definition_key not in schema:
+        return new_schema
+    else:
+        new_definitions = schema[definition_key].copy()
+    # Loop over existing definitions
+    for name, definition in schema[definition_key].items():
+        if name == class_name:
+            for prop in definition["properties"]:
+                if "description" in new_definitions[name]["properties"][prop]:
+                    raise ValueError(
+                        f"Property {name}.{prop} already has description"
+                    )
+                else:
+                    new_definitions[name]["properties"][prop][
+                        "description"
+                    ] = descriptions[prop]
+    new_schema[definition_key] = new_definitions
+    logging.info("[_insert_class_attrs_descriptions] END")
+    return new_schema

fractal_task_tools/_package_name_tools.py

@@ -0,0 +1,27 @@
+import re
+
+
+def normalize_package_name(pkg_name: str) -> str:
+    """
+    Implement both PyPa and custom package-name normalization
+
+    1. PyPa normalization: The name should be lowercased with all runs of the
+        characters `.`, `-`, or `_` replaced with a single `-` character
+        (https://packaging.python.org/en/latest/specifications/name-normalization).
+    2. Custom normalization: Replace `-` with `_`, to obtain the
+        imported-module name.
+
+    Args:
+        pkg_name: The non-normalized package name.
+
+    Returns:
+        The normalized package name.
+    """
+
+    # Apply PyPa normalization
+    pypa_normalized_package_name = re.sub(r"[-_.]+", "-", pkg_name).lower()
+
+    # Replace `-` with `_`
+    final_package_name = pypa_normalized_package_name.replace("-", "_")
+
+    return final_package_name

fractal_task_tools/_pydantic_generatejsonschema.py

@@ -0,0 +1,81 @@
+"""
+Custom Pydantic v2 JSON Schema generation tools.
+
+As of Pydantic V2, the JSON Schema representation of model attributes marked
+as `Optional` changed, and the new behavior consists in marking the
+corresponding properties as an `anyOf` of either a `null` or the actual type.
+This is not always the required behavior, see e.g.
+* https://github.com/pydantic/pydantic/issues/7161
+* https://github.com/pydantic/pydantic/issues/8394
+
+Here we list some alternative ways of reverting this change.
+"""
+import logging
+
+from pydantic.json_schema import GenerateJsonSchema
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core.core_schema import WithDefaultSchema
+
+logger = logging.getLogger("CustomGenerateJsonSchema")
+
+
+class CustomGenerateJsonSchema(GenerateJsonSchema):
+    def get_flattened_anyof(
+        self, schemas: list[JsonSchemaValue]
+    ) -> JsonSchemaValue:
+        null_schema = {"type": "null"}
+        if null_schema in schemas:
+            logger.warning(
+                "Drop `null_schema` before calling `get_flattened_anyof`"
+            )
+            schemas.pop(schemas.index(null_schema))
+        return super().get_flattened_anyof(schemas)
+
+    def default_schema(self, schema: WithDefaultSchema) -> JsonSchemaValue:
+        json_schema = super().default_schema(schema)
+        if "default" in json_schema.keys() and json_schema["default"] is None:
+            logger.warning(f"Pop `None` default value from {json_schema=}")
+            json_schema.pop("default")
+        return json_schema
+
+
+# class GenerateJsonSchemaA(GenerateJsonSchema):
+#     def nullable_schema(self, schema):
+#         null_schema = {"type": "null"}
+#         inner_json_schema = self.generate_inner(schema["schema"])
+#         if inner_json_schema == null_schema:
+#             return null_schema
+#         else:
+#             logging.info("A: Skip calling `get_flattened_anyof` method")
+#             return inner_json_schema
+
+
+# class GenerateJsonSchemaB(GenerateJsonSchemaA):
+#     def default_schema(self, schema: WithDefaultSchema) -> JsonSchemaValue:
+#         original_json_schema = super().default_schema(schema)
+#         new_json_schema = deepcopy(original_json_schema)
+#         default = new_json_schema.get("default", None)
+#         if default is None:
+#             logging.info("B: Pop None default")
+#             new_json_schema.pop("default")
+#         return new_json_schema
+
+
+# class GenerateJsonSchemaC(GenerateJsonSchema):
+#     def get_flattened_anyof(
+#         self, schemas: list[JsonSchemaValue]
+#     ) -> JsonSchemaValue:
+
+#         original_json_schema_value = super().get_flattened_anyof(schemas)
+#         members = original_json_schema_value.get("anyOf")
+#         logging.info("C", original_json_schema_value)
+#         if (
+#             members is not None
+#             and len(members) == 2
+#             and {"type": "null"} in members
+#         ):
+#             new_json_schema_value = {"type": [t["type"] for t in members]}
+#             logging.info("C", new_json_schema_value)
+#             return new_json_schema_value
+#         else:
+#             return original_json_schema_value

fractal_task_tools/_signature_constraints.py

@@ -0,0 +1,143 @@
+import inspect
+import logging
+from importlib import import_module
+from inspect import Parameter
+from inspect import Signature
+from inspect import signature
+from pathlib import Path
+from typing import Any
+
+from ._union_types import is_annotated_union
+from ._union_types import is_tagged
+from ._union_types import is_union
+
+
+# The following variables are copied from `pydantic.v1.decorator`
+# (for pydantic v2.11.10)
+ALT_V_ARGS = "v__args"
+ALT_V_KWARGS = "v__kwargs"
+V_DUPLICATE_KWARGS = "v__duplicate_kwargs"
+V_POSITIONAL_ONLY_NAME = "v__positional_only"
+
+FORBIDDEN_PARAM_NAMES = (
+    "args",
+    "kwargs",
+    V_POSITIONAL_ONLY_NAME,
+    V_DUPLICATE_KWARGS,
+    ALT_V_ARGS,
+    ALT_V_KWARGS,
+)
+
+
+def _extract_function(
+    module_relative_path: str,
+    function_name: str,
+    package_name: str,
+    verbose: bool = False,
+) -> callable:
+    """
+    Extract function from a module with the same name.
+
+    Args:
+        package_name: Example `fractal_tasks_core`.
+        module_relative_path: Example `tasks/create_ome_zarr.py`.
+        function_name: Example `create_ome_zarr`.
+        verbose:
+    """
+    if not module_relative_path.endswith(".py"):
+        raise ValueError(f"{module_relative_path=} must end with '.py'")
+    module_relative_path_no_py = str(
+        Path(module_relative_path).with_suffix("")
+    )
+    module_relative_path_dots = module_relative_path_no_py.replace("/", ".")
+    if verbose:
+        logging.info(
+            f"Now calling `import_module` for "
+            f"{package_name}.{module_relative_path_dots}"
+        )
+    imported_module = import_module(
+        f"{package_name}.{module_relative_path_dots}"
+    )
+    if verbose:
+        logging.info(
+            f"Now getting attribute {function_name} from "
+            f"imported module {imported_module}."
+        )
+    task_function = getattr(imported_module, function_name)
+    return task_function
+
+
+def _validate_plain_union(
+    *,
+    param: Parameter,
+    _type: Any,
+) -> None:
+    """
+    Fail for known cases of invalid plain-union types.
+
+    A plain union annotation is (by construction) one for which
+    `is_union(_type) = True`. The only supported forms of plain unions
+    are `X | None` or `X | None = None` (or equivalent forms).
+
+    Note that `Optional[X]` is equivalent to `X | None` and thus it also gets
+    validated through this function.
+
+    Args:
+        param: The full `inspect.Parameter` object.
+        _type:
+            The type annotation to review. Note that this may be equal to
+            `param.annotation` or to `param.annotation.__origin__` (when the
+            original `param.annotation` is an `Annotated` union).
+    """
+    args = _type.__args__
+    if len(args) != 2:
+        raise ValueError(
+            "Only unions of two elements are supported, but parameter "
+            f"'{param.name}' has type hint '{_type}'."
+        )
+    elif not any(arg is type(None) for arg in args):
+        raise ValueError(
+            "One union element must be None, but parameter "
+            f"'{param.name}' has type hint '{_type}'."
+        )
+    elif (param.default is not None) and (param.default != inspect._empty):
+        raise ValueError(
+            "Non-None default not supported, but parameter "
+            f"'{param.name}' has type hint '{_type}' "
+            f"and default {param.default}."
+        )
+
+
+def _validate_function_signature(function: callable) -> Signature:
+    """
+    Validate the function signature of a task.
+
+    Implement a set of checks for type hints that do not play well with the
+    creation of JSON Schema, see issue 399 in `fractal-tasks-core` and issue
+    65 in `fractal-task-tools`.
+
+    Args:
+        function: A callable function.
+    """
+    sig = signature(function)
+    for param in sig.parameters.values():
+        # Check that name is not forbidden
+        if param.name in FORBIDDEN_PARAM_NAMES:
+            raise ValueError(
+                f"Function {function} has argument with forbidden "
+                "name '{{param.name}}'"
+            )
+        # Validate plain unions or non-tagged annotated unions
+        if is_union(param.annotation):
+            _validate_plain_union(
+                _type=param.annotation,
+                param=param,
+            )
+        elif is_annotated_union(param.annotation):
+            if not is_tagged(param.annotation):
+                _validate_plain_union(
+                    _type=param.annotation.__origin__,
+                    param=param,
+                )
+    logging.info("[_validate_function_signature] END")
+    return sig

fractal_task_tools/_task_arguments.py

@@ -0,0 +1,75 @@
+import logging
+from typing import Any
+from typing import Literal
+
+
+REQUIRED_ARGUMENTS: dict[tuple[str, str], set[str]] = {
+    ("non_parallel", "non_parallel"): {"zarr_urls", "zarr_dir"},
+    ("compound", "non_parallel"): {"zarr_urls", "zarr_dir"},
+    ("parallel", "parallel"): {"zarr_url"},
+    ("compound", "parallel"): {"zarr_url", "init_args"},
+    ("converter_non_parallel", "non_parallel"): {"zarr_dir"},
+    ("converter_compound", "non_parallel"): {"zarr_dir"},
+    ("converter_compound", "parallel"): {"zarr_url", "init_args"},
+}
+FORBIDDEN_ARGUMENTS: dict[tuple[str, str], set[str]] = {
+    ("non_parallel", "non_parallel"): {"zarr_url"},
+    ("compound", "non_parallel"): {"zarr_url"},
+    ("parallel", "parallel"): {"zarr_urls", "zarr_dir"},
+    ("compound", "parallel"): {"zarr_urls", "zarr_dir"},
+    ("converter_non_parallel", "non_parallel"): {"zarr_url", "zarr_urls"},
+    ("converter_compound", "non_parallel"): {"zarr_url", "zarr_urls"},
+    ("converter_compound", "parallel"): {"zarr_urls", "zarr_dir"},
+}
+
+
+def validate_arguments(
+    *,
+    task_type: Literal["parallel", "non_parallel", "compound"],
+    executable_kind: Literal["parallel", "non_parallel"],
+    schema: dict[str, Any],
+) -> None:
+    """
+    Validate schema arguments against required/forbidden ones.
+
+    Arguments:
+        task_type:
+        executable_kind: The `parallel`/`non_parallel` part of the task.
+        schema:
+    """
+
+    key = (task_type, executable_kind)
+    if not (key in REQUIRED_ARGUMENTS and key in FORBIDDEN_ARGUMENTS):
+        logging.error(f"Invalid {task_type=}, {executable_kind=}.")
+        raise ValueError(f"Invalid {task_type=}, {executable_kind=}.")
+
+    required_args = REQUIRED_ARGUMENTS[key]
+    forbidden_args = FORBIDDEN_ARGUMENTS[key]
+
+    schema_properties = set(schema["properties"].keys())
+
+    logging.info(
+        f"[validate_arguments] Task has arguments: {schema_properties}"
+    )
+    logging.info(f"[validate_arguments] Required arguments: {required_args}")
+    logging.info(f"[validate_arguments] Forbidden arguments: {forbidden_args}")
+
+    missing_required_arguments = {
+        arg for arg in required_args if arg not in schema_properties
+    }
+    if missing_required_arguments:
+        error_msg = (
+            "[validate_arguments] Required arguments "
+            f"{missing_required_arguments} are missing."
+        )
+        logging.error(error_msg)
+        raise ValueError(error_msg)
+
+    present_forbidden_args = forbidden_args.intersection(schema_properties)
+    if present_forbidden_args:
+        error_msg = (
+            "[validate_arguments] Forbidden arguments "
+            f"{present_forbidden_args} are present."
+        )
+        logging.error(error_msg)
+        raise ValueError(error_msg)

fractal_task_tools/_task_docs.py

@@ -0,0 +1,109 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from docstring_parser import parse as docparse
+
+from ._descriptions import _get_function_docstring
+
+
+def _get_function_description(
+    package_name: str, module_path: str, function_name: str
+) -> str:
+    """
+    Extract function description from its docstring.
+
+    Args:
+        package_name: Example `fractal_tasks_core`.
+        module_path: Example `tasks/create_ome_zarr.py`.
+        function_name: Example `create_ome_zarr`.
+    """
+    # Extract docstring from ast.FunctionDef
+    docstring = _get_function_docstring(
+        package_name=package_name,
+        module_path=module_path,
+        function_name=function_name,
+    )
+    # Parse docstring (via docstring_parser)
+    parsed_docstring = docparse(docstring)
+    # Combine short/long descriptions (if present)
+    short_description = parsed_docstring.short_description
+    long_description = parsed_docstring.long_description
+    items = []
+    if short_description:
+        items.append(short_description)
+    if long_description:
+        items.append(long_description)
+    if items:
+        if parsed_docstring.blank_after_short_description:
+            return "\n\n".join(items)
+        else:
+            return "\n".join(items)
+    else:
+        return ""
+
+
+def create_docs_info(
+    *,
+    executable_non_parallel: Optional[str] = None,
+    executable_parallel: Optional[str] = None,
+    package: str,
+) -> str:
+    """
+    Return task description based on function docstring.
+    """
+    logging.info("[create_docs_info] START")
+    docs_info = []
+    for executable in [executable_non_parallel, executable_parallel]:
+        if executable is None:
+            continue
+        # Extract the function name.
+        # Note: this could be made more general, but for the moment we assume
+        # that the function has the same name as the module)
+        function_name = Path(executable).with_suffix("").name
+        logging.info(f"[create_docs_info] {function_name=}")
+        # Get function description
+        description = _get_function_description(
+            package_name=package,
+            module_path=executable,
+            function_name=function_name,
+        )
+        docs_info.append(f"## {function_name}\n{description}\n")
+    docs_info = "".join(docs_info)
+    logging.info("[create_docs_info] END")
+    return docs_info
+
+
+def read_docs_info_from_file(
+    *,
+    docs_info: str,
+    task_list_path: str,
+) -> str:
+    """
+    Return task description based on the content of a file.
+
+    An example of valid argument is
+    ```
+    docs_info = "file:relative/path/info.md"
+    ```
+    where the path is relative to the folder where `task_list.py` is.
+    """
+    logging.info("[read_docs_info_from_file] START")
+
+    # Preliminary checks
+    if not docs_info.startswith("file:"):
+        raise ValueError(f"Invalid docs_info='{docs_info}'.")
+    relative_path = Path(docs_info[5:])
+    if relative_path.is_absolute():
+        raise ValueError(
+            f"Invalid docs_info='{docs_info}' (path must be relative)."
+        )
+
+    base_path = Path(task_list_path).parent
+    docs_path = (base_path / relative_path).as_posix()
+    logging.info(f"[read_docs_info_from_file] Reading docs from {docs_path}")
+    with open(docs_path, "r") as f:
+        docs_info = f.read()
+    logging.info("[read_docs_info_from_file] END")
+
+    return docs_info