From d462428418708f364e2c0f316edff48e8504e465 Mon Sep 17 00:00:00 2001 From: Christopher Neugebauer Date: Wed, 14 Sep 2022 08:59:36 -0700 Subject: [PATCH] Add context manager to handle temporary unfreezing of `frozen_after_init` objects # Rust tests and lints will be skipped. Delete if not intended. [ci skip-rust] # Building wheels and fs_util will be skipped. Delete if not intended. [ci skip-build-wheels] (cherry picked from commit a6190ecc6142b040c5013a809cbfb2f5c32cc80f) --- src/python/pants/util/meta.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/python/pants/util/meta.py b/src/python/pants/util/meta.py index 9ab047cb923..6c257b8271b 100644 --- a/src/python/pants/util/meta.py +++ b/src/python/pants/util/meta.py @@ -2,9 +2,10 @@ # Licensed under the Apache License, Version 2.0 (see LICENSE). from abc import ABC, abstractmethod +from contextlib import contextmanager from dataclasses import FrozenInstanceError as FrozenInstanceError from functools import wraps -from typing import Any, Callable, Optional, Type, TypeVar, Union +from typing import Any, Callable, Iterator, Optional, Type, TypeVar, Union T = TypeVar("T") C = TypeVar("C", bound=Type) @@ -125,6 +126,15 @@ def freeze_instance(self) -> None: def unfreeze_instance(self) -> None: self._is_frozen = False + @contextmanager + def unfrozen(self) -> Iterator: + old_is_frozen = self._is_frozen + try: + self._is_frozen = False + yield + finally: + self._is_frozen = old_is_frozen + @wraps(prev_init) def new_init(self, *args: Any, **kwargs: Any) -> None: prev_init(self, *args, **kwargs) @@ -140,6 +150,7 @@ def new_setattr(self, key: str, value: Any) -> None: cls._freeze_instance = freeze_instance cls._unfreeze_instance = unfreeze_instance + cls._unfrozen = unfrozen cls.__init__ = new_init cls.__setattr__ = new_setattr # type: ignore[assignment]