Python contextlib.contextmanager() Examples

def ensure_safe_environment_variables():
    """
    Get a context manager to safely set environment variables

    All changes will be undone on close, hence environment variables set
    within this contextmanager will neither persist nor change global state.
    """
    saved_environ = dict(os.environ)
    try:
        yield
    finally:
        os.environ.clear()
        os.environ.update(saved_environ)


# -----------------------------------------------------------------------------
# Comparators 

def testContextManager(self):
    in_context = [False for i in xrange(10)]

    @contextlib.contextmanager
    def enter_into_context(i):
      in_context[i] = True
      try:
        yield
      finally:
        in_context[i] = False

    parallelized_context = parallelizer.SyncParallelizer(
        [enter_into_context(i) for i in xrange(10)])

    with parallelized_context:
      self.assertTrue(all(in_context))
    self.assertFalse(any(in_context)) 

def test_custom_context_manager(self):
        """Test that @custom_contextmanager is recognized as configured."""
        node = astroid.extract_node(
            """
        from contextlib import contextmanager
        def custom_contextmanager(f):
            return contextmanager(f)
        @custom_contextmanager
        def dec():
            yield
        with dec():
            pass
        """
        )
        with self.assertNoMessages():
            self.checker.visit_with(node) 

def test_custom_context_manager(self):
        """Test that @custom_contextmanager is recognized as configured."""
        node = astroid.extract_node(
            """
        from contextlib import contextmanager
        def custom_contextmanager(f):
            return contextmanager(f)
        @custom_contextmanager
        def dec():
            yield
        with dec():
            pass
        """
        )
        with self.assertNoMessages():
            self.checker.visit_with(node) 

def test_position_of_obj_unwraps():
    import contextlib

    @contextlib.contextmanager
    def cm():
        raise NotImplementedError()

    pdb_ = PdbTest()
    pos = pdb_._get_position_of_obj(cm)

    if hasattr(inspect, "unwrap"):
        assert pos[0] == THIS_FILE_CANONICAL
        assert pos[2] == [
            "    @contextlib.contextmanager\n",
            "    def cm():\n",
            "        raise NotImplementedError()\n",
        ]
    else:
        contextlib_file = contextlib.__file__
        if sys.platform == 'win32':
            contextlib_file = contextlib_file.lower()
        assert pos[0] == contextlib_file.rstrip("c") 

def test_weak_values_destroy_while_iterating(self):
        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
        dict, objects = self.make_weak_valued_dict()
        self.check_weak_destroy_while_iterating(dict, objects, 'iterkeys')
        self.check_weak_destroy_while_iterating(dict, objects, 'iteritems')
        self.check_weak_destroy_while_iterating(dict, objects, 'itervalues')
        self.check_weak_destroy_while_iterating(dict, objects, 'itervaluerefs')
        dict, objects = self.make_weak_valued_dict()
        @contextlib.contextmanager
        def testcontext():
            try:
                it = iter(dict.iteritems())
                next(it)
                # Schedule a key/value for removal and recreate it
                k = objects.pop().arg
                gc.collect()      # just in case
                yield k, Object(k)
            finally:
                it = None           # should commit all removals
                gc.collect()
        self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext) 

def test_weak_keys_destroy_while_iterating(self):
        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
        dict, objects = self.make_weak_keyed_dict()
        self.check_weak_destroy_while_iterating(dict, objects, 'iterkeys')
        self.check_weak_destroy_while_iterating(dict, objects, 'iteritems')
        self.check_weak_destroy_while_iterating(dict, objects, 'itervalues')
        self.check_weak_destroy_while_iterating(dict, objects, 'iterkeyrefs')
        dict, objects = self.make_weak_keyed_dict()
        @contextlib.contextmanager
        def testcontext():
            try:
                it = iter(dict.iteritems())
                next(it)
                # Schedule a key/value for removal and recreate it
                v = objects.pop().arg
                gc.collect()      # just in case
                yield Object(v), v
            finally:
                it = None           # should commit all removals
                gc.collect()
        self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext) 

def run_xz_decompress(stdout, stderr=None):
    """Run xz --decompress and return a contextmanager object for the proc."""
    xz = None
    try:
        xz = subprocess.Popen(["xz", "--decompress", "--stdout"],
                              stdin=subprocess.PIPE, stdout=stdout,
                              stderr=stderr)
        yield xz
    finally:
        if not xz:
            raise OSError("You must have an 'xz' binary in PATH.")

        xz.stdin.close()
        result = xz.wait()

        if result != 0:
            raise IOError("xz --decompress returned %d." % result) 

def setUpClass(
        cls,
        launch_service,
        proc_info,
        proc_output
    ):
        @contextlib.contextmanager
        def launch_pkg_command(self, arguments, **kwargs):
            pkg_command_action = ExecuteProcess(
                cmd=['ros2', 'pkg', *arguments],
                additional_env={'PYTHONUNBUFFERED': '1'},
                name='ros2pkg-cli',
                output='screen',
                **kwargs
            )
            with launch_testing.tools.launch_process(
                launch_service, pkg_command_action, proc_info, proc_output
            ) as pkg_command:
                yield pkg_command
        cls.launch_pkg_command = launch_pkg_command 

def solve_generator(
    *, call: Callable, stack: AsyncExitStack, sub_values: Dict[str, Any]
) -> Any:
    if is_gen_callable(call):
        cm = contextmanager_in_threadpool(contextmanager(call)(**sub_values))
    elif is_async_gen_callable(call):
        if not inspect.isasyncgenfunction(call):
            # asynccontextmanager from the async_generator backfill pre python3.7
            # does not support callables that are not functions or methods.
            # See https://github.com/python-trio/async_generator/issues/32
            #
            # Expand the callable class into its __call__ method before decorating it.
            # This approach will work on newer python versions as well.
            call = getattr(call, "__call__", None)
        cm = asynccontextmanager(call)(**sub_values)
    return await stack.enter_async_context(cm) 

def _patch_AuthAPILdapDataBasedMethodTestMixIn():
    @contextlib.contextmanager
    def monkey_patched_standard_context(self, search_flat_return_value):
        create_and_init_db(timeout=20)
        try:
            populate_db_with_test_data(self.__class__.__name__)
            with self._singleton_off():
                self.auth_api = AuthAPI(settings=prepare_auth_db_settings())
                try:
                    with patch.object(AuthAPI, '_get_root_node',
                                      AuthAPI._get_root_node.func):  # unmemoized (not cached)
                        yield
                finally:
                    self.auth_api = None
        finally:
            drop_db()

    def monkey_patched_assert_problematic_orgs_logged(self, *args, **kwargs):
        pass

    test_auth_api._AuthAPILdapDataBasedMethodTestMixIn.standard_context = \
        monkey_patched_standard_context
    test_auth_api._AuthAPILdapDataBasedMethodTestMixIn.assert_problematic_orgs_logged = \
        monkey_patched_assert_problematic_orgs_logged 

def forced_piece_size(pytestconfig):
    @contextlib.contextmanager
    def _forced_piece_size(piece_size):
        orig_piece_size_min = torf.Torrent.piece_size_min
        torf.Torrent.piece_size_min = piece_size
        with mock.patch('torf.Torrent.piece_size', new_callable=mock.PropertyMock) as mock_piece_size:
            def piece_size_setter(prop, torrent, value):
                torrent.metainfo['info']['piece length'] = piece_size

            mock_piece_size.return_value = piece_size
            mock_piece_size.__set__ = piece_size_setter
            yield piece_size
        torf.Torrent.piece_size_min = orig_piece_size_min
    return _forced_piece_size


# https://stackoverflow.com/a/45690594 

def modify_config(request):
    @contextlib.contextmanager
    def _modify_config(f, d):
        """ Overwrites yaml file ``f`` with values in ``dict`` ``d`` """
        orig = yaml.load(open(f, 'r'))
        modified = orig.copy()
        try:
            modified = deep_update(modified, d)
            tmpcfg = tempfile.mkstemp(prefix='yatsm_', suffix='.yaml')[1]
            yaml.dump(modified, open(tmpcfg, 'w'))
            yield tmpcfg
        except:
            raise
        finally:
            os.remove(tmpcfg)
    return _modify_config


# RASTER READING UTILS 

def open(self, mode: QIODevice.OpenMode) -> contextlib.closing:
        """Open the underlying device and ensure opening succeeded.

        Raises OSError if opening failed.

        Args:
            mode: QIODevice::OpenMode flags.

        Return:
            A contextlib.closing() object so this can be used as
            contextmanager.
        """
        ok = self.dev.open(mode)
        if not ok:
            raise QtOSError(self.dev)
        return contextlib.closing(self) 

def setUpClass(
        cls,
        launch_service,
        proc_info,
        proc_output,
        rmw_implementation
    ):
        @contextlib.contextmanager
        def launch_action_command(self, arguments):
            action_command_action = ExecuteProcess(
                cmd=['ros2', 'action', *arguments],
                name='ros2action-cli', output='screen',
                additional_env={
                    'RMW_IMPLEMENTATION': rmw_implementation,
                    'PYTHONUNBUFFERED': '1'
                }
            )
            with launch_testing.tools.launch_process(
                launch_service, action_command_action, proc_info, proc_output,
                output_filter=launch_testing_ros.tools.basic_output_filter(
                    filtered_rmw_implementation=rmw_implementation
                )
            ) as action_command:
                yield action_command
        cls.launch_action_command = launch_action_command 

def test_weak_keys_destroy_while_iterating(self):
        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
        dict, objects = self.make_weak_keyed_dict()
        self.check_weak_destroy_while_iterating(dict, objects, 'iterkeys')
        self.check_weak_destroy_while_iterating(dict, objects, 'iteritems')
        self.check_weak_destroy_while_iterating(dict, objects, 'itervalues')
        self.check_weak_destroy_while_iterating(dict, objects, 'iterkeyrefs')
        dict, objects = self.make_weak_keyed_dict()
        @contextlib.contextmanager
        def testcontext():
            try:
                it = iter(dict.iteritems())
                next(it)
                # Schedule a key/value for removal and recreate it
                v = objects.pop().arg
                gc.collect()      # just in case
                yield Object(v), v
            finally:
                it = None           # should commit all removals
                gc.collect()
        self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext) 

def download_miniconda_installer(installer_url, sha256sum):
    """
    Context manager to download miniconda installer from a given URL

    This should be used as a contextmanager. It downloads miniconda installer
    of given version, verifies the sha256sum & provides path to it to the `with`
    block to run.
    """
    with tempfile.NamedTemporaryFile() as f:
        with open(f.name, 'wb') as f:
            f.write(requests.get(installer_url).content)

        if sha256_file(f.name) != sha256sum:
            raise Exception('sha256sum hash mismatch! Downloaded file corrupted')

        yield f.name 

def test_weak_values_destroy_while_iterating(self):
        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
        dict, objects = self.make_weak_valued_dict()
        self.check_weak_destroy_while_iterating(dict, objects, 'iterkeys')
        self.check_weak_destroy_while_iterating(dict, objects, 'iteritems')
        self.check_weak_destroy_while_iterating(dict, objects, 'itervalues')
        self.check_weak_destroy_while_iterating(dict, objects, 'itervaluerefs')
        dict, objects = self.make_weak_valued_dict()
        @contextlib.contextmanager
        def testcontext():
            try:
                it = iter(dict.iteritems())
                next(it)
                # Schedule a key/value for removal and recreate it
                k = objects.pop().arg
                gc.collect()      # just in case
                yield k, Object(k)
            finally:
                it = None           # should commit all removals
                gc.collect()
        self.check_weak_destroy_and_mutate_while_iterating(dict, testcontext) 

def maybe_raises(exception=None):
  @contextmanager
  def noop():
    yield

  with (noop() if exception is None else pytest.raises(exception)):
    yield 

def saving(self, *args):
        '''
        Context manager to facilitate writing the movie file.

        ``*args`` are any parameters that should be passed to `setup`.
        '''
        # This particular sequence is what contextlib.contextmanager wants
        self.setup(*args)
        yield
        self.finish() 

def _build_config_create_interface_submode_context(self, configurations):
        @contextlib.contextmanager
        def multiple_submode_context():
            with configurations.submode_context('interface {}'.format(self.parent_interface.name), cancel_empty=True):
                with configurations.submode_context('service instance {} ethernet'.format(self.service_instance)):
                    yield

        return multiple_submode_context() 

def _setup_outputs(root_output_dir, experiment_name, hparam_dict):
  """Set up directories for experiment loops, write hyperparameters to disk."""

  if not experiment_name:
    raise ValueError('experiment_name must be specified.')

  create_if_not_exists(root_output_dir)

  checkpoint_dir = os.path.join(root_output_dir, 'checkpoints', experiment_name)
  create_if_not_exists(checkpoint_dir)
  checkpoint_mngr = checkpoint_manager.FileCheckpointManager(checkpoint_dir)

  results_dir = os.path.join(root_output_dir, 'results', experiment_name)
  create_if_not_exists(results_dir)
  metrics_mngr = metrics_manager.ScalarMetricsManager(
      results_dir, use_bz2=FLAGS.write_metrics_with_bz2)

  summary_logdir = os.path.join(root_output_dir, 'logdir', experiment_name)
  create_if_not_exists(summary_logdir)
  summary_writer = tf.summary.create_file_writer(summary_logdir)

  hparam_dict['metrics_file'] = metrics_mngr.metrics_filename
  hparams_file = os.path.join(results_dir, 'hparams.csv')
  utils_impl.atomic_write_to_csv(pd.Series(hparam_dict), hparams_file)

  logging.info('Writing...')
  logging.info('    checkpoints to: %s', checkpoint_dir)
  logging.info('    metrics csv to: %s', metrics_mngr.metrics_filename)
  logging.info('    summaries to: %s', summary_logdir)

  @contextlib.contextmanager
  def profiler(round_num):
    if (FLAGS.rounds_per_profile > 0 and
        round_num % FLAGS.rounds_per_profile == 0):
      with tf.profiler.experimental.Profile(summary_logdir):
        yield
    else:
      yield

  return checkpoint_mngr, metrics_mngr, summary_writer, profiler 

def test_weak_destroy_and_mutate_while_iterating(self):
        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
        items = [SomeClass(c) for c in string.ascii_letters]
        s = WeakSet(items)
        @contextlib.contextmanager
        def testcontext():
            try:
                it = iter(s)
                next(it)
                # Schedule an item for removal and recreate it
                u = SomeClass(str(items.pop()))
                gc.collect()      # just in case
                yield u
            finally:
                it = None           # should commit all removals

        with testcontext() as u:
            self.assertNotIn(u, s)
        with testcontext() as u:
            self.assertRaises(KeyError, s.remove, u)
        self.assertNotIn(u, s)
        with testcontext() as u:
            s.add(u)
        self.assertIn(u, s)
        t = s.copy()
        with testcontext() as u:
            s.update(t)
        self.assertEqual(len(s), len(t))
        with testcontext() as u:
            s.clear()
        self.assertEqual(len(s), 0) 

def named_context(self, name):
        @contextlib.contextmanager
        def context():
            self.named_contexts.append(name)
            try:
                yield
            finally:
                self.assertEqual(self.named_contexts.pop(), name)
        return context 

def named_context(self, name):
        @contextlib.contextmanager
        def context():
            self.named_contexts.append(name)
            try:
                yield
            finally:
                self.assertEqual(self.named_contexts.pop(), name)
        return context 

def assertRaises(_exception, _callable=None, *args, **kwargs):
    """assertRaises that is usable as context manager or in a with statement

    Exceptions that don't match the given Exception type fall through::

    >>> with assertRaises(ValueError):
    ...     raise TypeError("banana")
    ...
    Traceback (most recent call last):
        ...
    TypeError: banana

    If it raises the given Exception type, the test passes
    >>> with assertRaises(KeyError):
    ...     dct = dict()
    ...     dct["apple"]

    If the expected error doesn't occur, it raises an error.
    >>> with assertRaises(KeyError):
    ...     dct = {'apple':True}
    ...     dct["apple"]
    Traceback (most recent call last):
        ...
    AssertionError: KeyError not raised.

    In addition to using it as a contextmanager, you can also use it as a
    function, just like the normal assertRaises

    >>> assertRaises(TypeError, ",".join, [1, 3, 5]);
    """
    manager = _AssertRaisesContextmanager(exception=_exception)
    # don't return anything if used in function form
    if _callable is not None:
        with manager:
            _callable(*args, **kwargs)
    else:
        return manager 

def set_trace():
    from IPython.core.debugger import Pdb
    try:
        Pdb(color_scheme='Linux').set_trace(sys._getframe().f_back)
    except:
        from pdb import Pdb as OldPdb
        OldPdb().set_trace(sys._getframe().f_back)

#------------------------------------------------------------------------------
# contextmanager to ensure the file cleanup 

def saving(self, *args):
        '''
        Context manager to facilitate writing the movie file.

        ``*args`` are any parameters that should be passed to `setup`.
        '''
        # This particular sequence is what contextlib.contextmanager wants
        self.setup(*args)
        yield
        self.finish() 

def test_weak_destroy_and_mutate_while_iterating(self):
        # Issue #7105: iterators shouldn't crash when a key is implicitly removed
        items = [SomeClass(c) for c in string.ascii_letters]
        s = WeakSet(items)
        @contextlib.contextmanager
        def testcontext():
            try:
                it = iter(s)
                next(it)
                # Schedule an item for removal and recreate it
                u = SomeClass(str(items.pop()))
                gc.collect()      # just in case
                yield u
            finally:
                it = None           # should commit all removals

        with testcontext() as u:
            self.assertNotIn(u, s)
        with testcontext() as u:
            self.assertRaises(KeyError, s.remove, u)
        self.assertNotIn(u, s)
        with testcontext() as u:
            s.add(u)
        self.assertIn(u, s)
        t = s.copy()
        with testcontext() as u:
            s.update(t)
        self.assertEqual(len(s), len(t))
        with testcontext() as u:
            s.clear()
        self.assertEqual(len(s), 0) 

def _load_reward_net_as_fn(shaped: bool) -> RewardFnLoaderFn:
    @contextlib.contextmanager
    def loader(path: str, venv: VecEnv) -> Iterator[common.RewardFn]:
        """Load train (shaped) or test (not shaped) reward from path."""
        del venv  # Unused.
        with networks.make_session() as (graph, sess):
            net = reward_net.RewardNet.load(path)
            reward = net.reward_output_train if shaped else net.reward_output_test

            def rew_fn(
                obs: np.ndarray,
                act: np.ndarray,
                next_obs: np.ndarray,
                dones: np.ndarray,
            ) -> np.ndarray:
                fd = {
                    net.obs_ph: obs,
                    net.act_ph: act,
                    net.next_obs_ph: next_obs,
                    net.done_ph: dones,
                }
                rew = sess.run(reward, feed_dict=fd)
                assert rew.shape == (len(obs),)
                return rew

            yield rew_fn

    return loader