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 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 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 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 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_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 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 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 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 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 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 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 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 _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 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 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 stdoutIO():
    """Stdout redirection as a context manager to evaluating code mutations."""

    @contextlib.contextmanager
    def stdoutIO(stdout=None):
        old = sys.stdout
        if stdout is None:
            stdout = StringIO()
        sys.stdout = stdout
        yield stdout
        sys.stdout = old

    return stdoutIO 

def NoIdleAssertion(reason):
  """Context manager for creating and releasing a NoIdleAssertion.

  https://docs.python.org/2/library/contextlib.html#contextlib.contextmanager

  Usage:
  with NoIdleAssertion():
    # Some stuff

  Args:
    reason: string, tag for the power assertion
  Yields:
    None
  """
  assertion_type = 'NoIdleSleepAssertion'
  io_lib = ConfigureIOKit()
  returncode, assertion_id = CreatePowerAssertion(
      io_lib, assertion_type, 255, reason)
  if returncode:
    logging.error('Could not create assertion: %s', returncode)
  else:
    logging.debug('Created %s', assertion_type)

  try:
    yield
  finally:
    returncode = ReleasePowerAssertion(io_lib, assertion_id)
    if returncode:
      logging.error('Could not release assertion: %s', returncode)
    else:
      logging.debug('Released %s', assertion_type) 

def noclose(file):
    """
    A "no-op" contextmanager that prevents files from closing.
    """
    try:
        yield file
    finally:
        pass 

def parent_step(self, name_tokens):
    """Opens a parent step.

    Returns a contextmanager object yielding (StepPresentation, List[StepData]).
    Refer to RecipeEngine.parent_step for details.
    """
    return self._engine.parent_step(name_tokens) 

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 set_trace():
    from IPython.core.debugger import Pdb
    try:
        Pdb(color_scheme='Linux').set_trace(sys._getframe().f_back)
    except Exception:
        from pdb import Pdb as OldPdb
        OldPdb().set_trace(sys._getframe().f_back)

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

def __enter__(self):
        """ For contextmanager-style usage. """
        self.start()
        return self 

def __exit__(self, exc_type, value, traceback):
        """ For contextmanager-style usage. """
        self.stop() 

def acquire_driver_object(self, uid, allocator):
        """Return a context manager to acquire a driver object

        Example
        -------
        >>> with back_ds.acquire(uid) as gdal_obj:
        ...     pass
        """
        @contextlib.contextmanager
        def _acquire():
            with self._ap_lock:
                if uid in self._ap_idle:
                    obj = self._ap_idle.pop_first_occurrence(uid)
                    allocate = False
                else:
                    self._ensure_one_slot()
                    allocate = True
                self._ap_used[uid] += 1

            if allocate:
                try:
                    obj = allocator()
                except:
                    with self._ap_lock:
                        self._ap_used[uid] -= 1
                    raise

            try:
                yield obj
            finally:
                with self._ap_lock:
                    self._ap_used[uid] -= 1
                    assert self._ap_used[uid] >= 0
                    self._ap_idle.push_front(uid, obj)
                    if self._ap_used[uid] == 0:
                        del self._ap_used[uid]

        return _acquire() 

def contextmanager(target):
  """A tf_decorator-aware wrapper for `contextlib.contextmanager`.

  Usage is identical to `contextlib.contextmanager`.

  Args:
    target: A callable to be wrapped in a contextmanager.
  Returns:
    A callable that can be used inside of a `with` statement.
  """
  context_manager = _contextlib.contextmanager(target)
  return tf_decorator.make_decorator(target, context_manager, 'contextmanager') 

def random_seed():
    @contextlib.contextmanager
    def _random_seed(seed):
        random.seed(seed)
        yield
        random.seed()
    return _random_seed