Added typing information for all public APIs (#1158)

* Added typing information for all public APIs

The docstrings were also slightly more standardized.

Only public facing APIs had the type information added to them.

There are a few minor bug fixes in the default values for some functions
(namely using [] as a default value (replaced with None)) as well
as removing default values where it does not make sense (S3.info() for example).
There should be no other functional change.

* Fix typos

* More clean-ups

* More cleanups

* Remove circular dep

* Fixups and cleanups

* Fix typo; make parameters in Parameter explicit

* Forgot an import

* Proper defaults for init for Parameter

* Fix issue with Parameter moving away from kwargs

* Re-added optional `key` value in S3.get()

* Addressed comments

* Proper type for  in Parameter

* Remove typeing for external packages

* Addressed comments

* Update docs for secrets

* Forgot two files

metaflow/current.py

@@ -1,5 +1,6 @@
 from collections import namedtuple
 import os
+from typing import Any, Optional
 
 Parallel = namedtuple("Parallel", ["main_ip", "num_nodes", "node_index"])
 
@@ -57,14 +58,14 @@ def _update_env(self, env):
         for k, v in env.items():
             setattr(self.__class__, k, property(fget=lambda _, v=v: v))
 
-    def __contains__(self, key):
+    def __contains__(self, key: str):
         return getattr(self, key, None) is not None
 
-    def get(self, key, default=None):
+    def get(self, key: str, default=None) -> Optional[Any]:
         return getattr(self, key, default)
 
     @property
-    def is_running_flow(self):
+    def is_running_flow(self) -> bool:
         """
         Returns True if called inside a running Flow, False otherwise.
 
@@ -79,7 +80,7 @@ def is_running_flow(self):
         return self._is_running
 
     @property
-    def flow_name(self):
+    def flow_name(self) -> str:
         """
         The name of the currently executing flow.
 
@@ -91,7 +92,7 @@ def flow_name(self):
         return self._flow_name
 
     @property
-    def run_id(self):
+    def run_id(self) -> str:
         """
         The run ID of the currently executing run.
 
@@ -103,7 +104,7 @@ def run_id(self):
         return self._run_id
 
     @property
-    def step_name(self):
+    def step_name(self) -> str:
         """
         The name of the currently executing step.
 
@@ -115,7 +116,7 @@ def step_name(self):
         return self._step_name
 
     @property
-    def task_id(self):
+    def task_id(self) -> str:
         """
         The task ID of the currently executing task.
 
@@ -127,7 +128,7 @@ def task_id(self):
         return self._task_id
 
     @property
-    def retry_count(self):
+    def retry_count(self) -> int:
         """
         The index of the task execution attempt.
 
@@ -144,7 +145,7 @@ def retry_count(self):
         return self._retry_count
 
     @property
-    def origin_run_id(self):
+    def origin_run_id(self) -> Optional[str]:
         """
         The run ID of the original run this run was resumed from.
 
@@ -157,13 +158,13 @@ def origin_run_id(self):
 
         Returns
         -------
-        str
+        str, optional
             Run ID of the original run.
         """
         return self._origin_run_id
 
     @property
-    def pathspec(self):
+    def pathspec(self) -> str:
         """
         Pathspec of the current run, i.e. a unique
         identifier of the current task. The returned
@@ -189,7 +190,7 @@ def pathspec(self):
         return "/".join(pathspec_components)
 
     @property
-    def namespace(self):
+    def namespace(self) -> str:
         """
         The current namespace.
 
@@ -201,13 +202,13 @@ def namespace(self):
         return self._namespace
 
     @property
-    def username(self):
+    def username(self) -> Optional[str]:
         """
         The name of the user who started the run, if available.
 
         Returns
         -------
-        str
+        str, optional
             User name.
         """
         return self._username

metaflow/flowspec.py

@@ -5,6 +5,7 @@
 
 from itertools import islice
 from types import FunctionType, MethodType
+from typing import Any, Callable, List, Optional, Tuple
 
 from . import cmd_with_io
 from .parameters import DelayedEvaluationParameter, Parameter
@@ -23,6 +24,9 @@
     basestring = str
 
 
+from .datastore.inputs import Inputs
+
+
 class InvalidNextException(MetaflowException):
     headline = "Invalid self.next() transition detected"
 
@@ -83,7 +87,7 @@ def __init__(self, use_cli=True):
 
         Parameters
         ----------
-        use_cli : bool, optional, default: True
+        use_cli : bool, default: True
             Set to True if the flow is invoked from __main__ or the command line
         """
 
@@ -104,7 +108,7 @@ def __init__(self, use_cli=True):
             cli.main(self)
 
     @property
-    def script_name(self):
+    def script_name(self) -> str:
         """
         [Legacy function - do not use. Use `current` instead]
 
@@ -214,7 +218,7 @@ def __iter__(self):
         """
         return iter(self._steps)
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str):
         if self._datastore and name in self._datastore:
             # load the attribute from the datastore...
             x = self._datastore[name]
@@ -231,7 +235,7 @@ def cmd(self, cmdline, input={}, output=[]):
         return cmd_with_io.cmd(cmdline, input=input, output=output)
 
     @property
-    def index(self):
+    def index(self) -> Optional[int]:
         """
         The index of this foreach branch.
 
@@ -244,14 +248,14 @@ def index(self):
 
         Returns
         -------
-        int
+        int, optional
             Index of the task in a foreach step.
         """
         if self._foreach_stack:
             return self._foreach_stack[-1].index
 
     @property
-    def input(self):
+    def input(self) -> Optional[Any]:
         """
         The value of the foreach artifact in this foreach branch.
 
@@ -264,12 +268,12 @@ def input(self):
 
         Returns
         -------
-        object
+        object, optional
             Input passed to the foreach task.
         """
         return self._find_input()
 
-    def foreach_stack(self):
+    def foreach_stack(self) -> Optional[List[Tuple[int, int, Any]]]:
         """
         Returns the current stack of foreach indexes and values for the current step.
 
@@ -353,7 +357,12 @@ def _find_input(self, stack_index=None):
                     )
             return self._cached_input[stack_index]
 
-    def merge_artifacts(self, inputs, exclude=[], include=[]):
+    def merge_artifacts(
+        self,
+        inputs: Inputs,
+        exclude: Optional[List[str]] = None,
+        include: Optional[List[str]] = None,
+    ) -> None:
         """
         Helper function for merging artifacts in a join step.
 
@@ -386,7 +395,7 @@ def merge_artifacts(self, inputs, exclude=[], include=[]):
 
         Parameters
         ----------
-        inputs : List[Steps]
+        inputs : Inputs
             Incoming steps to the join point.
         exclude : List[str], optional
             If specified, do not consider merging artifacts with a name in `exclude`.
@@ -405,6 +414,8 @@ def merge_artifacts(self, inputs, exclude=[], include=[]):
             This exception is thrown in case an artifact specified in `include` cannot
             be found.
         """
+        include = include or []
+        exclude = exclude or []
         node = self._graph[self._current_step]
         if node.type != "join":
             msg = (
@@ -488,7 +499,7 @@ def _validate_ubf_step(self, step_name):
             )
             raise InvalidNextException(msg)
 
-    def next(self, *dsts, **kwargs):
+    def next(self, *dsts: Callable[..., None], **kwargs) -> None:
         """
         Indicates the next step to execute after this step has completed.
 

metaflow/includefile.py

@@ -6,6 +6,7 @@
 import os
 
 from hashlib import sha1
+from typing import Any, Dict, Optional
 
 from metaflow._vendor import click
 
@@ -64,7 +65,7 @@ class IncludedFile(object):
     # the actual content)
 
     # @tracefunc
-    def __init__(self, descriptor):
+    def __init__(self, descriptor: Dict[str, Any]):
         self._descriptor = descriptor
         self._cached_size = None
 
@@ -245,22 +246,28 @@ class IncludeFile(Parameter):
     default : str or a function
         Default path to a local file. A function
         implies that the parameter corresponds to a *deploy-time parameter*.
-    is_text : bool
-        Convert the file contents to a string using the provided `encoding` (default: True).
+    is_text : bool, default: True
+        Convert the file contents to a string using the provided `encoding`.
         If False, the artifact is stored in `bytes`.
-    encoding : str
-        Use this encoding to decode the file contexts if `is_text=True` (default: `utf-8`).
-    required : bool
+    encoding : str, optional, default: 'utf-8'
+        Use this encoding to decode the file contexts if `is_text=True`.
+    required : bool, default: False
         Require that the user specified a value for the parameter.
         `required=True` implies that the `default` is not used.
-    help : str
+    help : str, optional
         Help text to show in `run --help`.
-    show_default : bool
-        If True, show the default value in the help text (default: True).
+    show_default : bool, default: True
+        If True, show the default value in the help text.
     """
 
     def __init__(
-        self, name, required=False, is_text=True, encoding=None, help=None, **kwargs
+        self,
+        name: str,
+        required: bool = False,
+        is_text: bool = True,
+        encoding: str = "utf-8",
+        help: Optional[str] = None,
+        **kwargs: Dict[str, str]
     ):
         # If a default is specified, it needs to be uploaded when the flow is deployed
         # (for example when doing a `step-functions create`) so we make the default

metaflow/parameters.py

@@ -1,5 +1,5 @@
 import json
-from collections import namedtuple
+from typing import Any, Callable, Dict, NamedTuple, Optional, Type, Union
 
 from metaflow._vendor import click
 
@@ -20,14 +20,14 @@
 # ParameterContext allows deploy-time functions modify their
 # behavior based on the context. We can add fields here without
 # breaking backwards compatibility but don't remove any fields!
-ParameterContext = namedtuple(
+ParameterContext = NamedTuple(
     "ParameterContext",
     [
-        "flow_name",
-        "user_name",
-        "parameter_name",
-        "logger",
-        "ds_type",
+        ("flow_name", str),
+        ("user_name", str),
+        ("parameter_name", str),
+        ("logger", Callable[..., None]),
+        ("ds_type", str),
     ],
 )
 
@@ -246,21 +246,52 @@ class MyFlow(FlowSpec):
         indicate that the value must be a valid JSON object. A function
         implies that the parameter corresponds to a *deploy-time parameter*.
         The type of the default value is used as the parameter `type`.
-    type : type
+    type : Type, default: None
         If `default` is not specified, define the parameter type. Specify
-        one of `str`, `float`, `int`, `bool`, or `JSONType` (default: str).
-    help : str
+        one of `str`, `float`, `int`, `bool`, or `JSONType`. If None, defaults
+        to the type of `default` or `str` if none specified.
+    help : str, optional
         Help text to show in `run --help`.
-    required : bool
+    required : bool, default: False
         Require that the user specified a value for the parameter.
         `required=True` implies that the `default` is not used.
-    show_default : bool
-        If True, show the default value in the help text (default: True).
+    show_default : bool, default: True
+        If True, show the default value in the help text.
     """
 
-    def __init__(self, name, **kwargs):
+    def __init__(
+        self,
+        name: str,
+        default: Optional[
+            Union[
+                str,
+                float,
+                int,
+                bool,
+                Dict[str, Any],
+                Callable[[], Union[str, float, int, bool, Dict[str, Any]]],
+            ]
+        ] = None,
+        type: Optional[
+            Union[Type[str], Type[float], Type[int], Type[bool], JSONTypeClass]
+        ] = None,
+        help: Optional[str] = None,
+        required: bool = False,
+        show_default: bool = True,
+        **kwargs: Dict[str, Any]
+    ):
         self.name = name
         self.kwargs = kwargs
+        for k, v in {
+            "default": default,
+            "type": type,
+            "help": help,
+            "required": required,
+            "show_default": show_default,
+        }.items():
+            if v is not None:
+                self.kwargs[k] = v
+
         # TODO: check that the type is one of the supported types
         param_type = self.kwargs["type"] = self._get_type(kwargs)
         reserved_params = [
@@ -289,8 +320,6 @@ def __init__(self, name, **kwargs):
                     "have a function as its value" % (name, field)
                 )
 
-        self.kwargs["show_default"] = self.kwargs.get("show_default", True)
-
         # default can be defined as a function
         default_field = self.kwargs.get("default")
         if callable(default_field) and not isinstance(default_field, DeployTimeField):