import cloudpickle

from contextlib import contextmanager


@contextmanager
def _no_globals():
  """
  Random variables hanging around in our namespace (like imported modules) get included when we try to pickle
  functions. If any of those globals happen to be generators or other unsupported types, that call will fail. This
  context manager provides a simple way of removing those variables so we can easily serialize pure functions.
  """
  previous_globals = globals().copy()
  try:
    for key in previous_globals.keys():
      del globals()[key]
    yield
  finally:
    for key, value in previous_globals.items():
      globals()[key] = value


def dumps_fn(fn):
  with _no_globals():
    import cloudpickle
    return cloudpickle.dumps(fn)


def loads_fn(serialized):
  return cloudpickle.loads(serialized)


_FN_KEY_ATTRS = [ 'co_code', 'co_consts', 'co_names' ]
def build_fn_key(fn):
  return cloudpickle.dumps(tuple( getattr(fn.__code__, i) for i in _FN_KEY_ATTRS ))


def test_index_fn(fn):
  '''
  Function serialization can be difficult to understand and the underlying mechanism is subject to change. This
  function allows you to test that your function still performs reliably when serialized and deserialized.

  Returns the resolved function that you can run to make absolutely sure the behavior is consistent
  '''

  key = build_fn_key(fn)
  resolved = loads_fn(dumps_fn(fn))
  resolved_key = build_fn_key(resolved)
  assert key == resolved_key, 'Function key does not match when serialized and deserialized'
  return resolved
