def freeze()

in antlir/btrfs_diff/freeze.py [0:0]

• 35 lines of code
• 14 McCabe index (conditional complexity)

def freeze(obj, *, _memo=None, **kwargs):
    # Don't bother memoizing primitive types
    if isinstance(obj, (bytes, Enum, float, int, str, type(None))):
        return obj

    if _memo is None:
        _memo = {}

    if id(obj) in _memo:  # Already frozen?
        return _memo[id(obj)]

    if hasattr(obj, "freeze"):
        frozen = obj.freeze(_memo=_memo, **kwargs)
    else:
        # At the moment, I don't have a need for passing extra data into
        # items that live inside containers.  If we're relaxing this, just
        # be sure to add `**kwargs` to each `freeze()` call below.
        assert kwargs == {}, kwargs
        # This is a lame-o way of identifying `NamedTuple`s. Using
        # `deepfrozen` would avoid this kludge.
        if (
            isinstance(obj, tuple)
            and hasattr(obj, "_replace")
            and hasattr(obj, "_fields")
            and hasattr(obj, "_make")
        ):
            frozen = obj._make(freeze(i, _memo=_memo) for i in obj)
        elif isinstance(obj, (list, tuple)):
            frozen = tuple(freeze(i, _memo=_memo) for i in obj)
        elif isinstance(obj, dict):
            frozen = frozendict(
                {
                    freeze(k, _memo=_memo): freeze(v, _memo=_memo)
                    for k, v in obj.items()
                }
            )
        elif isinstance(obj, (set, frozenset)):
            frozen = frozenset(freeze(i, _memo=_memo) for i in obj)
        elif isinstance(obj, DoNotFreeze):
            frozen = obj
        else:
            raise NotImplementedError(type(obj))

    _memo[id(obj)] = frozen
    return frozen