pytorch/torch/fx/immutable_collections.py
James Reed 538647fe1f [WIP][FX] BC guarantees for 1.10 (#63888)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63888

Test Plan: Imported from OSS

Reviewed By: pbelevich

Differential Revision: D30523133

Pulled By: jamesr66a

fbshipit-source-id: b04cc0d842a74862f42ecba98b757310cd2ec7b0
2021-08-30 19:56:46 -07:00

29 lines
1.3 KiB
Python

from ._compatibility import compatibility
_help_mutation = """\
If you are attempting to modify the kwargs or args of a torch.fx.Node object,
instead create a new copy of it and assign the copy to the node:
new_args = ... # copy and mutate args
node.args = new_args
"""
def _no_mutation(self, *args, **kwargs):
raise NotImplementedError(f"'{type(self).__name__}' object does not support mutation. {_help_mutation}")
def _create_immutable_container(base, mutable_functions):
container = type('immutable_' + base.__name__, (base,), {})
for attr in mutable_functions:
setattr(container, attr, _no_mutation)
return container
immutable_list = _create_immutable_container(list,
['__delitem__', '__iadd__', '__imul__', '__setitem__', 'append',
'clear', 'extend', 'insert', 'pop', 'remove'])
immutable_list.__reduce__ = lambda self: (immutable_list, (tuple(iter(self)),))
compatibility(is_backward_compatible=True)(immutable_list)
immutable_dict = _create_immutable_container(dict, ['__delitem__', '__setitem__', 'clear', 'pop', 'popitem', 'update'])
immutable_dict.__reduce__ = lambda self: (immutable_dict, (iter(self.items()),))
compatibility(is_backward_compatible=True)(immutable_dict)