text: from functools import partial
  import re
  import os
  import sys
  
  from .flowspec import FlowSpec
  from .exception import (
      MetaflowInternalError,
      MetaflowException,
      InvalidDecoratorAttribute,
  )
  
  import click
  
  
  class BadStepDecoratorException(MetaflowException):
      headline = "Syntax error"
  
      def __init__(self, deco, func):
          msg = (
              "You tried to apply decorator '{deco}' on '{func}' which is "
              "not declared as a @step. Make sure you apply this decorator "
              "on a function which has @step on the line just before the "
              "function name and @{deco} is above @step.".format(
                  deco=deco, func=func.__name__
              )
          )
          super(BadStepDecoratorException, self).__init__(msg)
  
  
  class BadFlowDecoratorException(MetaflowException):
      headline = "Syntax error"
  
      def __init__(self, deconame):
          msg = (
              "Decorator '%s' can be applied only to FlowSpecs. Make sure "
              "the decorator is above a class definition." % deconame
          )
          super(BadFlowDecoratorException, self).__init__(msg)
  
  
  class UnknownStepDecoratorException(MetaflowException):
      headline = "Unknown step decorator"
  
      def __init__(self, deconame):
          from .plugins import STEP_DECORATORS
  
          decos = ", ".join(
              t.name for t in STEP_DECORATORS if not t.name.endswith("_internal")
          )
          msg = (
              "Unknown step decorator *{deconame}*. The following decorators are "
              "supported: *{decos}*".format(deconame=deconame, decos=decos)
          )
          super(UnknownStepDecoratorException, self).__init__(msg)
  
  
  class DuplicateStepDecoratorException(MetaflowException):
      headline = "Duplicate decorators"
  
      def __init__(self, deco, func):
          msg = (
              "Step '{step}' already has a decorator '{deco}'. "
              "You can specify each decorator only once.".format(
                  step=func.__name__, deco=deco
              )
          )
          super(DuplicateStepDecoratorException, self).__init__(msg)
  
  
  class UnknownFlowDecoratorException(MetaflowException):
      headline = "Unknown flow decorator"
  
      def __init__(self, deconame):
          from .plugins import FLOW_DECORATORS
  
          decos = ", ".join(t.name for t in FLOW_DECORATORS)
          msg = (
              "Unknown flow decorator *{deconame}*. The following decorators are "
              "supported: *{decos}*".format(deconame=deconame, decos=decos)
          )
          super(UnknownFlowDecoratorException, self).__init__(msg)
  
  
  class DuplicateFlowDecoratorException(MetaflowException):
      headline = "Duplicate decorators"
  
      def __init__(self, deco):
          msg = (
              "Flow already has a decorator '{deco}'. "
              "You can specify each decorator only once.".format(deco=deco)
          )
          super(DuplicateFlowDecoratorException, self).__init__(msg)
  
  
  class Decorator(object):
      """
      Base class for all decorators.
      """
  
      name = "NONAME"
      defaults = {}
  
      def __init__(self, attributes=None, statically_defined=False):
          self.attributes = self.defaults.copy()
          self.statically_defined = statically_defined
  
          if attributes:
              for k, v in attributes.items():
                  if k in self.defaults:
                      self.attributes[k] = v
                  else:
                      raise InvalidDecoratorAttribute(self.name, k, self.defaults)
  
      @classmethod
      def _parse_decorator_spec(cls, deco_spec):
          top = deco_spec.split(":", 1)
          if len(top) == 1:
              return cls()
          else:
              name, attrspec = top
              attrs = dict(
                  map(lambda x: x.strip(), a.split("="))
                  for a in re.split(""",(?=[\s\w]+=)""", attrspec.strip("\"'"))
              )
              return cls(attributes=attrs)
  
      def make_decorator_spec(self):
          attrs = {k: v for k, v in self.attributes.items() if v is not None}
          if attrs:
              attrstr = ",".join("%s=%s" % x for x in attrs.items())
              return "%s:%s" % (self.name, attrstr)
          else:
              return self.name
  
      def __str__(self):
          mode = "decorated" if self.statically_defined else "cli"
          attrs = " ".join("%s=%s" % x for x in self.attributes.items())
          if attrs:
              attrs = " " + attrs
          fmt = "%s<%s%s>" % (self.name, mode, attrs)
          return fmt
  
  
  class FlowDecorator(Decorator):
  
      _flow_decorators = []
      options = {}
  
      def __init__(self, *args, **kwargs):
          # Note that this assumes we are executing one flow per process so we have a global list of
          # _flow_decorators. A similar setup is used in parameters.
          self._flow_decorators.append(self)
          super(FlowDecorator, self).__init__(*args, **kwargs)
  
      def flow_init(
          self, flow, graph, environment, flow_datastore, metadata, logger, echo, options
      ):
          """
          Called when all decorators have been created for this flow.
          """
          pass
  
      def get_top_level_options(self):
          """
          Return a list of option-value pairs that correspond to top-level
          options that should be passed to subprocesses (tasks). The option
          names should be a subset of the keys in self.options.
  
          If the decorator has a non-empty set of options in self.options, you
          probably want to return the assigned values in this method.
          """
          return []
  
  
  # compare this to parameters.add_custom_parameters
  def add_decorator_options(cmd):
      seen = {}
      for deco in flow_decorators():
          for option, kwargs in deco.options.items():
              if option in seen:
                  msg = (
                      "Flow decorator '%s' uses an option '%s' which is also "
                      "used by the decorator '%s'. This is a bug in Metaflow. "
                      "Please file a ticket on GitHub."
                      % (deco.name, option, seen[option])
                  )
                  raise MetaflowInternalError(msg)
              else:
                  seen[option] = deco.name
                  cmd.params.insert(0, click.Option(("--" + option,), **kwargs))
      return cmd
  
  
  def flow_decorators():
      return FlowDecorator._flow_decorators
  
  
  class StepDecorator(Decorator):
      """
      Base class for all step decorators.
  
      Example:
  
      @my_decorator
      @step
      def a(self):
          pass
  
      @my_decorator
      @step
      def b(self):
          pass
  
      To make the above work, define a subclass
  
      class MyDecorator(StepDecorator):
          name = "my_decorator"
  
      and include it in plugins.STEP_DECORATORS. Now both a() and b()
      get an instance of MyDecorator, so you can keep step-specific
      state easily.
  
      TODO (savin): Initialize the decorators with flow, graph,
                    step.__name__ etc., so that we don't have to
                    pass them around with every lifecycle call.
      """
  
      def step_init(
          self, flow, graph, step_name, decorators, environment, flow_datastore, logger
      ):
          """
          Called when all decorators have been created for this step
          """
          pass
  
      def package_init(self, flow, step_name, environment):
          """
          Called to determine package components
          """
          pass
  
      def add_to_package(self):
          """
          Called to add custom packages needed for a decorator. This hook will be
          called in the `MetaflowPackage` class where metaflow compiles the code package
          tarball. This hook is invoked in the `MetaflowPackage`'s `path_tuples`
          function. The `path_tuples` function is a generator that yields a tuple of
          `(file_path, arcname)`.`file_path` is the path of the file in the local file system;
          the `arcname` is the path of the file in the constructed tarball or the path of the file
          after decompressing the tarball.
  
          Returns a list of tuples where each tuple represents (file_path, arcname)
          """
          return []
  
      def step_task_retry_count(self):
          """
          Called to determine the number of times this task should be retried.
          Returns a tuple of (user_code_retries, error_retries). Error retries
          are attempts to run the process after the user code has failed all
          its retries.
  
          Typically, the runtime takes the maximum of retry counts across
          decorators and user specification to determine the task retry count.
          If you want to force no retries, return the special values (None, None).
          """
          return 0, 0
  
      def runtime_init(self, flow, graph, package, run_id):
          """
          Top-level initialization before anything gets run in the runtime
          context.
          """
          pass
  
      def runtime_task_created(
          self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
      ):
          """
          Called when the runtime has created a task related to this step.
          """
          pass
  
      def runtime_finished(self, exception):
          """
          Called when the runtime created task finishes or encounters an interrupt/exception.
          """
          pass
  
      def runtime_step_cli(
          self, cli_args, retry_count, max_user_code_retries, ubf_context
      ):
          """
          Access the command line for a step execution in the runtime context.
          """
          pass
  
      def task_pre_step(
          self,
          step_name,
          task_datastore,
          metadata,
          run_id,
          task_id,
          flow,
          graph,
          retry_count,
          max_user_code_retries,
          ubf_context,
          inputs,
      ):
          """
          Run before the step function in the task context.
          """
          pass
  
      def task_decorate(
          self, step_func, flow, graph, retry_count, max_user_code_retries, ubf_context
      ):
          return step_func
  
      def task_post_step(
          self, step_name, flow, graph, retry_count, max_user_code_retries
      ):
          """
          Run after the step function has finished successfully in the task
          context.
          """
          pass
  
      def task_exception(
          self, exception, step_name, flow, graph, retry_count, max_user_code_retries
      ):
          """
          Run if the step function raised an exception in the task context.
  
          If this method returns True, it is assumed that the exception has
          been taken care of and the flow may continue.
          """
          pass
  
      def task_finished(
          self, step_name, flow, graph, is_task_ok, retry_count, max_user_code_retries
      ):
          """
          Run after the task context has been finalized.
  
          is_task_ok is set to False if the user code raised an exception that
          was not handled by any decorator.
  
          Note that you can't create or modify data artifacts in this method
          since the task has been finalized by the time this method
          is called. Also note that the task may fail after this method has been
          called, so this method may get called multiple times for a task over
          multiple attempts, similar to all task_ methods.
          """
          pass
  
  
  def _base_flow_decorator(decofunc, *args, **kwargs):
      """
      Decorator prototype for all flow (class) decorators. This function gets
      specialized and imported for all decorators types by
      _import_plugin_decorators().
      """
      if args:
          # No keyword arguments specified for the decorator, e.g. @foobar.
          # The first argument is the class to be decorated.
          cls = args[0]
          if isinstance(cls, type) and issubclass(cls, FlowSpec):
              # flow decorators add attributes in the class dictionary,
              # _flow_decorators.
              if decofunc.name in cls._flow_decorators:
                  raise DuplicateFlowDecoratorException(decofunc.name)
              else:
                  cls._flow_decorators[decofunc.name] = decofunc(
                      attributes=kwargs, statically_defined=True
                  )
          else:
              raise BadFlowDecoratorException(decofunc.name)
          return cls
      else:
          # Keyword arguments specified, e.g. @foobar(a=1, b=2).
          # Return a decorator function that will get the actual
          # function to be decorated as the first argument.
          def wrap(f):
              return _base_flow_decorator(decofunc, f, **kwargs)
  
          return wrap
  
  
  def _base_step_decorator(decotype, *args, **kwargs):
      """
      Decorator prototype for all step decorators. This function gets specialized
      and imported for all decorators types by _import_plugin_decorators().
      """
      if args:
          # No keyword arguments specified for the decorator, e.g. @foobar.
          # The first argument is the function to be decorated.
          func = args[0]
          if not hasattr(func, "is_step"):
              raise BadStepDecoratorException(decotype.name, func)
  
          # Only the first decorator applies
          if decotype.name in [deco.name for deco in func.decorators]:
              raise DuplicateStepDecoratorException(decotype.name, func)
          else:
              func.decorators.append(decotype(attributes=kwargs, statically_defined=True))
  
          return func
      else:
          # Keyword arguments specified, e.g. @foobar(a=1, b=2).
          # Return a decorator function that will get the actual
          # function to be decorated as the first argument.
          def wrap(f):
              return _base_step_decorator(decotype, f, **kwargs)
  
          return wrap
  
  
  def _attach_decorators(flow, decospecs):
      """
      Attach decorators to all steps during runtime. This has the same
      effect as if you defined the decorators statically in the source for
      every step. Used by --with command line parameter.
      """
      # Attach the decorator to all steps that don't have this decorator
      # already. This means that statically defined decorators are always
      # preferred over runtime decorators.
      #
      # Note that each step gets its own instance of the decorator class,
      # so decorator can maintain step-specific state.
      for step in flow:
          _attach_decorators_to_step(step, decospecs)
  
  
  def _attach_decorators_to_step(step, decospecs):
      """
      Attach decorators to a step during runtime. This has the same
      effect as if you defined the decorators statically in the source for
      the step.
      """
      from .plugins import STEP_DECORATORS
  
      decos = {decotype.name: decotype for decotype in STEP_DECORATORS}
      for decospec in decospecs:
          deconame = decospec.strip("'").split(":")[0]
          if deconame not in decos:
              raise UnknownStepDecoratorException(deconame)
          # Attach the decorator to step if it doesn't have the decorator
          # already. This means that statically defined decorators are always
          # preferred over runtime decorators.
          if deconame not in [deco.name for deco in step.decorators]:
              deco = decos[deconame]._parse_decorator_spec(decospec)
              step.decorators.append(deco)
  
  
  def _init_flow_decorators(
      flow, graph, environment, flow_datastore, metadata, logger, echo, deco_options
  ):
      for deco in flow._flow_decorators.values():
          opts = {option: deco_options[option] for option in deco.options}
          deco.flow_init(
              flow, graph, environment, flow_datastore, metadata, logger, echo, opts
          )
  
  
  def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
      for step in flow:
          for deco in step.decorators:
              deco.step_init(
                  flow,
                  graph,
                  step.__name__,
                  step.decorators,
                  environment,
                  flow_datastore,
                  logger,
              )
  
  
  def step(f):
      """
      The step decorator. Makes a method a step in the workflow.
      """
      f.is_step = True
      f.decorators = []
      try:
          # python 3
          f.name = f.__name__
      except:
          # python 2
          f.name = f.__func__.func_name
      return f
  
  
  def _import_plugin_decorators(globals_dict):
      """
      Auto-generate a decorator function for every decorator
      defined in plugins.STEP_DECORATORS and plugins.FLOW_DECORATORS.
      """
      from .plugins import STEP_DECORATORS, FLOW_DECORATORS
  
      # Q: Why not use StepDecorators directly as decorators?
      # A: Getting an object behave as a decorator that can work
      #    both with and without arguments is surprisingly hard.
      #    It is easier to make plain function decorators work in
      #    the dual mode - see _base_step_decorator above.
      for decotype in STEP_DECORATORS:
          globals_dict[decotype.name] = partial(_base_step_decorator, decotype)
  
      # add flow-level decorators
      for decotype in FLOW_DECORATORS:
          globals_dict[decotype.name] = partial(_base_flow_decorator, decotype)
  