Skip to content

Commit

Permalink
[internal] Add context manager to handle temporary unfreezing of `fro…
Browse files Browse the repository at this point in the history
…zen_after_init` objects (#16863)

This adds a safer way to temporarily unfreeze `frozen_after_init` objects, for example in tests. The API is not exposed to static typing (just like the existing `_unfreeze_instance`), but it will re-freeze once finished.
  • Loading branch information
Christopher Neugebauer authored Sep 14, 2022
1 parent ebcf1a1 commit 7e266c7
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion src/python/pants/util/meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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]

Expand Down

0 comments on commit 7e266c7

Please sign in to comment.