mypyc: enable subclassing and/or pickling of popular classes

mr-c · mr-c · commit c5e291286344 · 2023-03-02T15:55:32.000+01:00
Fixes: DataBiosphere/toil#4394
diff --git a/cwltool/builder.py b/cwltool/builder.py
@@ -19,15 +19,15 @@
 
 from cwl_utils import expression
 from cwl_utils.file_formats import check_format
+from mypy_extensions import mypyc_attr
 from rdflib import Graph
+from ruamel.yaml.comments import CommentedMap
 from schema_salad.avro.schema import Names, Schema, make_avsc_object
 from schema_salad.exceptions import ValidationException
 from schema_salad.sourceline import SourceLine
 from schema_salad.utils import convert_to_dict, json_dumps
 from schema_salad.validate import validate
 
-from ruamel.yaml.comments import CommentedMap
-
 from .errors import WorkflowException
 from .loghandler import _logger
 from .mutation import MutationManager
@@ -94,6 +94,7 @@ def substitute(value: str, replace: str) -> str:
     return value + replace
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class Builder(HasReqsHints):
     """Helper class to construct a command line from a CWL CommandLineTool."""
 
diff --git a/cwltool/command_line_tool.py b/cwltool/command_line_tool.py
@@ -32,15 +32,15 @@
 )
 
 import shellescape
+from mypy_extensions import mypyc_attr
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.avro.schema import Schema
 from schema_salad.exceptions import ValidationException
 from schema_salad.ref_resolver import file_uri, uri_file_path
 from schema_salad.sourceline import SourceLine
 from schema_salad.utils import json_dumps
 from schema_salad.validate import validate_ex
 
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-
 from .builder import (
     INPUT_OBJ_VOCAB,
     Builder,
@@ -208,6 +208,7 @@ def run(
                 self.output_callback({}, "permanentFail")
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class ExpressionTool(Process):
     def job(
         self,
@@ -315,6 +316,7 @@ def revmap_file(builder: Builder, outdir: str, f: CWLObjectType) -> Optional[CWL
     )
 
 
+@mypyc_attr(serializable=True)
 class CallbackJob:
     """Callback Job class, used by :py:func:`CommandLineTool.job`."""
 
@@ -401,6 +403,7 @@ def __init__(self, msg: str, port: CWLObjectType, **kwargs: Any) -> None:
         )
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class CommandLineTool(Process):
     def __init__(self, toolpath_object: CommentedMap, loadingContext: LoadingContext) -> None:
         """Initialize this CommandLineTool."""
@@ -470,8 +473,8 @@ def make_job_runner(self, runtimeContext: RuntimeContext) -> Type[JobBase]:
             )
         return CommandLineJob
 
+    @staticmethod
     def make_path_mapper(
-        self,
         reffiles: List[CWLObjectType],
         stagedir: str,
         runtimeContext: RuntimeContext,
diff --git a/cwltool/context.py b/cwltool/context.py
@@ -18,13 +18,12 @@
     Union,
 )
 
+from ruamel.yaml.comments import CommentedMap
 from schema_salad.avro.schema import Names
 from schema_salad.ref_resolver import Loader
 from schema_salad.utils import FetcherCallableType
 from typing_extensions import Literal
 
-from ruamel.yaml.comments import CommentedMap
-
 from .mpi import MpiConfig
 from .pathmapper import PathMapper
 from .stdfsaccess import StdFsAccess
diff --git a/cwltool/cwlrdf.py b/cwltool/cwlrdf.py
@@ -4,11 +4,10 @@
 
 from rdflib import Graph
 from rdflib.query import ResultRow
+from ruamel.yaml.comments import CommentedMap
 from schema_salad.jsonld_context import makerdf
 from schema_salad.utils import ContextType
 
-from ruamel.yaml.comments import CommentedMap
-
 from .cwlviewer import CWLViewer
 from .process import Process
 
diff --git a/cwltool/docker.py b/cwltool/docker.py
@@ -83,7 +83,7 @@ def __init__(
         self,
         builder: Builder,
         joborder: CWLObjectType,
-        make_path_mapper: Callable[..., PathMapper],
+        make_path_mapper: Callable[[List[CWLObjectType], str, RuntimeContext, bool], PathMapper],
         requirements: List[CWLObjectType],
         hints: List[CWLObjectType],
         name: str,
@@ -452,7 +452,7 @@ def __init__(
         self,
         builder: Builder,
         joborder: CWLObjectType,
-        make_path_mapper: Callable[..., PathMapper],
+        make_path_mapper: Callable[[List[CWLObjectType], str, RuntimeContext, bool], PathMapper],
         requirements: List[CWLObjectType],
         hints: List[CWLObjectType],
         name: str,
diff --git a/cwltool/executors.py b/cwltool/executors.py
@@ -20,6 +20,7 @@
 )
 
 import psutil
+from mypy_extensions import mypyc_attr
 from schema_salad.exceptions import ValidationException
 from schema_salad.sourceline import SourceLine
 
@@ -40,6 +41,7 @@
 TMPDIR_LOCK = Lock()
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class JobExecutor(metaclass=ABCMeta):
     """Abstract base job executor."""
 
@@ -178,6 +180,7 @@ def check_for_abstract_op(tool: CWLObjectType) -> None:
         return (None, "permanentFail")
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class SingleJobExecutor(JobExecutor):
     """Default single-threaded CWL reference executor."""
 
diff --git a/cwltool/job.py b/cwltool/job.py
@@ -118,7 +118,7 @@ def __init__(
         self,
         builder: Builder,
         joborder: CWLObjectType,
-        make_path_mapper: Callable[..., PathMapper],
+        make_path_mapper: Callable[[List[CWLObjectType], str, RuntimeContext, bool], PathMapper],
         requirements: List[CWLObjectType],
         hints: List[CWLObjectType],
         name: str,
diff --git a/cwltool/load_tool.py b/cwltool/load_tool.py
@@ -21,6 +21,7 @@
 )
 
 from cwl_utils.parser import cwl_v1_2, cwl_v1_2_utils
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.exceptions import ValidationException
 from schema_salad.fetcher import Fetcher
 from schema_salad.ref_resolver import Loader, file_uri
@@ -34,8 +35,6 @@
     json_dumps,
 )
 
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-
 from . import CWL_CONTENT_TYPES, process, update
 from .context import LoadingContext
 from .errors import GraphTargetMissingException
diff --git a/cwltool/main.py b/cwltool/main.py
@@ -34,6 +34,9 @@
 import argcomplete
 import coloredlogs
 import pkg_resources  # part of setuptools
+import ruamel.yaml
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
+from ruamel.yaml.main import YAML
 from schema_salad.exceptions import ValidationException
 from schema_salad.ref_resolver import Loader, file_uri, uri_file_path
 from schema_salad.sourceline import cmap, strip_dup_lineno
@@ -45,10 +48,6 @@
     yaml_no_ts,
 )
 
-import ruamel.yaml
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-from ruamel.yaml.main import YAML
-
 from . import CWL_CONTENT_TYPES, workflow
 from .argparser import arg_parser, generate_parser, get_default_args
 from .context import LoadingContext, RuntimeContext, getdefault
diff --git a/cwltool/pack.py b/cwltool/pack.py
@@ -14,11 +14,10 @@
     cast,
 )
 
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.ref_resolver import Loader, SubLoader
 from schema_salad.utils import ResolveType
 
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-
 from .context import LoadingContext
 from .load_tool import fetch_document, resolve_and_validate_document
 from .process import shortname, uniquename
diff --git a/cwltool/pathmapper.py b/cwltool/pathmapper.py
@@ -7,6 +7,7 @@
 from pathlib import Path
 from typing import Dict, Iterator, List, Optional, Tuple, cast
 
+from mypy_extensions import mypyc_attr
 from schema_salad.exceptions import ValidationException
 from schema_salad.ref_resolver import uri_file_path
 from schema_salad.sourceline import SourceLine
@@ -42,6 +43,7 @@
 """
 
 
+@mypyc_attr(allow_interpreted_subclasses=True)
 class PathMapper:
     """
     Mapping of files from relative path provided in the file to a tuple.
@@ -192,11 +194,12 @@ def setup(self, referenced_files: List[CWLObjectType], basedir: str) -> None:
         for fob in referenced_files:
             if self.separateDirs:
                 stagedir = os.path.join(self.stagedir, "stg%s" % uuid.uuid4())
+            copy = cast(bool, fob.get("writable", False) or False)
             self.visit(
                 fob,
                 stagedir,
                 basedir,
-                copy=cast(bool, fob.get("writable", False)),
+                copy=copy,
                 staged=True,
             )
 
diff --git a/cwltool/process.py b/cwltool/process.py
@@ -36,6 +36,7 @@
 from mypy_extensions import mypyc_attr
 from pkg_resources import resource_stream
 from rdflib import Graph
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.avro.schema import (
     Names,
     Schema,
@@ -49,8 +50,6 @@
 from schema_salad.utils import convert_to_dict
 from schema_salad.validate import avro_type_name, validate_ex
 
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-
 from .builder import INPUT_OBJ_VOCAB, Builder
 from .context import LoadingContext, RuntimeContext, getdefault
 from .errors import UnsupportedRequirement, WorkflowException
diff --git a/cwltool/procgenerator.py b/cwltool/procgenerator.py
@@ -1,11 +1,10 @@
 import copy
 from typing import Dict, Optional, Tuple, cast
 
+from ruamel.yaml.comments import CommentedMap
 from schema_salad.exceptions import ValidationException
 from schema_salad.sourceline import indent
 
-from ruamel.yaml.comments import CommentedMap
-
 from .context import LoadingContext, RuntimeContext
 from .errors import WorkflowException
 from .load_tool import load_tool
diff --git a/cwltool/singularity.py b/cwltool/singularity.py
@@ -121,7 +121,7 @@ def __init__(
         self,
         builder: Builder,
         joborder: CWLObjectType,
-        make_path_mapper: Callable[..., PathMapper],
+        make_path_mapper: Callable[[List[CWLObjectType], str, RuntimeContext, bool], PathMapper],
         requirements: List[CWLObjectType],
         hints: List[CWLObjectType],
         name: str,
diff --git a/cwltool/update.py b/cwltool/update.py
@@ -11,12 +11,11 @@
     cast,
 )
 
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.exceptions import ValidationException
 from schema_salad.ref_resolver import Loader
 from schema_salad.sourceline import SourceLine
 
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-
 from .loghandler import _logger
 from .utils import CWLObjectType, CWLOutputType, aslist, visit_class, visit_field
 
diff --git a/cwltool/validate_js.py b/cwltool/validate_js.py
@@ -19,6 +19,7 @@
 from cwl_utils.expression import scanner as scan_expression
 from cwl_utils.sandboxjs import code_fragment_to_js, exec_js_process
 from pkg_resources import resource_stream
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.avro.schema import (
     ArraySchema,
     EnumSchema,
@@ -30,8 +31,6 @@
 from schema_salad.utils import json_dumps
 from schema_salad.validate import validate_ex
 
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
-
 from .errors import WorkflowException
 from .loghandler import _logger
 
diff --git a/cwltool/workflow.py b/cwltool/workflow.py
@@ -17,11 +17,10 @@
 )
 from uuid import UUID
 
+from ruamel.yaml.comments import CommentedMap
 from schema_salad.exceptions import ValidationException
 from schema_salad.sourceline import SourceLine, indent
 
-from ruamel.yaml.comments import CommentedMap
-
 from . import command_line_tool, context, procgenerator
 from .checker import circular_dependency_checker, loop_checker, static_checker
 from .context import LoadingContext, RuntimeContext, getdefault
diff --git a/setup.py b/setup.py
@@ -64,7 +64,7 @@
         "cwltool/main.py",
         "cwltool/mutation.py",
         "cwltool/pack.py",
-        # "cwltool/pathmapper.py",  # class PathMapper needs to be subclassable
+        "cwltool/pathmapper.py",
         "cwltool/process.py",
         "cwltool/procgenerator.py",
         # "cwltool/provenance.py",  # WritableBag is having issues
diff --git a/tests/test_anon_types.py b/tests/test_anon_types.py
@@ -1,11 +1,11 @@
 from typing import cast
 
 import pytest
+from ruamel.yaml.comments import CommentedMap
 from schema_salad.sourceline import cmap
 
 from cwltool.command_line_tool import CommandLineTool
 from cwltool.context import LoadingContext
-from ruamel.yaml.comments import CommentedMap
 
 snippet = cast(
     CommentedMap,
diff --git a/tests/test_cuda.py b/tests/test_cuda.py
@@ -5,13 +5,13 @@
 from schema_salad.avro import schema
 
 from cwltool.builder import Builder
+from cwltool.command_line_tool import CommandLineTool
 from cwltool.context import LoadingContext, RuntimeContext
 from cwltool.cuda import cuda_version_and_device_count
 from cwltool.errors import WorkflowException
 from cwltool.job import CommandLineJob
 from cwltool.load_tool import load_tool
 from cwltool.main import main
-from cwltool.pathmapper import PathMapper
 from cwltool.process import use_custom_schema
 from cwltool.stdfsaccess import StdFsAccess
 from cwltool.update import INTERNAL_VERSION
@@ -108,7 +108,7 @@ def test_cuda_job_setup_check(makedirs: MagicMock, check_output: MagicMock) -> N
 </nvidia>
 """
 
-    jb = CommandLineJob(builder, {}, PathMapper, [], [], "")
+    jb = CommandLineJob(builder, {}, CommandLineTool.make_path_mapper, [], [], "")
     jb._setup(runtime_context)
 
 
@@ -130,7 +130,7 @@ def test_cuda_job_setup_check_err(makedirs: MagicMock, check_output: MagicMock)
 <cuda_version>1.0</cuda_version>
 </nvidia>
 """
-    jb = CommandLineJob(builder, {}, PathMapper, [], [], "")
+    jb = CommandLineJob(builder, {}, CommandLineTool.make_path_mapper, [], [], "")
     with pytest.raises(WorkflowException):
         jb._setup(runtime_context)
 
diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -15,6 +15,7 @@
 import pytest
 from cwl_utils.errors import JavascriptException
 from cwl_utils.sandboxjs import param_re
+from ruamel.yaml.comments import CommentedMap, CommentedSeq
 from schema_salad.exceptions import ValidationException
 
 import cwltool.checker
@@ -28,7 +29,6 @@
 from cwltool.main import main
 from cwltool.process import CWL_IANA
 from cwltool.utils import CWLObjectType, dedup
-from ruamel.yaml.comments import CommentedMap, CommentedSeq
 
 from .util import get_data, get_main_output, needs_docker, working_directory
 
diff --git a/tests/test_mpi.py b/tests/test_mpi.py
diff --git a/tests/test_path_checks.py b/tests/test_path_checks.py
diff --git a/tests/test_provenance.py b/tests/test_provenance.py
diff --git a/tests/test_streaming.py b/tests/test_streaming.py
diff --git a/tests/test_subclass_mypyc.py b/tests/test_subclass_mypyc.py
diff --git a/tests/test_tmpdir.py b/tests/test_tmpdir.py

Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@`
`21`	`21`	`)`
`22`	`22`
`23`	`23`	`from cwl_utils.parser import cwl_v1_2, cwl_v1_2_utils`
	`24`	`+from ruamel.yaml.comments import CommentedMap, CommentedSeq`
`24`	`25`	`from schema_salad.exceptions import ValidationException`
`25`	`26`	`from schema_salad.fetcher import Fetcher`
`26`	`27`	`from schema_salad.ref_resolver import Loader, file_uri`
`@@ -34,8 +35,6 @@`
`34`	`35`	`json_dumps,`
`35`	`36`	`)`
`36`	`37`
`37`		`-from ruamel.yaml.comments import CommentedMap, CommentedSeq`
`38`		`-`
`39`	`38`	`from . import CWL_CONTENT_TYPES, process, update`
`40`	`39`	`from .context import LoadingContext`
`41`	`40`	`from .errors import GraphTargetMissingException`