Source code for fmeTestSupport.Context

# Copyright (c) Fraunhofer MEVIS, Germany. All rights reserved.
# **InsertLicense** code

from contextlib import contextmanager
import os
from typing import Callable, ContextManager, Dict, Union, TypeVar, MutableMapping

YieldType = TypeVar("YieldType")
MemoryType = TypeVar("MemoryType")

@contextmanager
def Context( doAction:Union[ Callable[ [], YieldType],
                             Callable[ [MemoryType], YieldType] ],
             undoAction:Union[ Callable[ [YieldType], None],
                               Callable[ [YieldType, MemoryType], None] ],
             undoOnError:bool=True,
             memory:MemoryType=None ) -> ContextManager[YieldType]:
  """Generic context manager to temporarily do something, and then undo it on exit, but (optionally) only if no exception has occurred in between.

  :param doAction: Callback to be called before yielding
  :param undoAction: Callback to be called on exit
  :param undoOnError: If set, the undoAction is NOT executed if an exception occurred within the context manager body
  :param memory: Optional mutable object (e.g. a dict) that doAction can use to remember stuff required by undoAction. \
    If provided, it is passed to both functions as first/second parameter.
  :yields: The return value of doAction
  """
  undo = True
  v = doAction() if memory is None else doAction( memory )
  try:
    yield v
  except:
    undo = undoOnError
    raise
  finally:
    if undo:
      undoAction( v ) if memory is None else undoAction( v, memory )

@contextmanager
def TemporaryItems( dictionary:MutableMapping, additionalItems:dict, unsetOnError:bool=True ):
  """
  Contextmanager that temporarily extends a dictionary with additional items. If the provided
  additionalItems overwrite pre-existing ones, the original valus will be restored afterwards.  

  :param dictionary: Target dictionary to be extended.
  :param additionalItems: Source dictionary to be used to extend the given target dictionary.
  :param unsetOnError: If set, the variables are NOT unset if an exception occurred within the context manager body.
  :yields: None
  """
  orig_values = {}
  def ___do():
    for variableName, value in additionalItems.items():
      if variableName in dictionary:
        orig_values[ variableName ] = dictionary[ variableName ]
      dictionary[ variableName ] = value
  def ___undo( _ ):
    for variableName in additionalItems:
      if variableName in orig_values:
        dictionary[ variableName ] = orig_values[ variableName ]
      else:
        del dictionary[ variableName ] 
  #
  with Context( ___do, ___undo, undoOnError=unsetOnError ):
    yield

@contextmanager
def EnvironmentVariables( additionalItems:Dict[str,str], unsetOnError:bool=True ):
  """
  Context manager to temporarily set environment variables. Internally uses AdditionalItems().

  :param additionalItems: Dictionary to be used to extend environment variables.
  :param unsetOnError: If set, the variables are NOT unset if an exception occurred within the context manager body.
  :yields: None
  """
  with TemporaryItems( os.environ, additionalItems, unsetOnError=unsetOnError ):
    yield