_mangling.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162
  1. """Import mangling.
  2. See mangling.md for details.
  3. """
  4. import re
  5. _mangle_index = 0
  6. class PackageMangler:
  7. """
  8. Used on import, to ensure that all modules imported have a shared mangle parent.
  9. """
  10. def __init__(self):
  11. global _mangle_index
  12. self._mangle_index = _mangle_index
  13. # Increment the global index
  14. _mangle_index += 1
  15. # Angle brackets are used so that there is almost no chance of
  16. # confusing this module for a real module. Plus, it is Python's
  17. # preferred way of denoting special modules.
  18. self._mangle_parent = f"<torch_package_{self._mangle_index}>"
  19. def mangle(self, name) -> str:
  20. assert len(name) != 0
  21. return self._mangle_parent + "." + name
  22. def demangle(self, mangled: str) -> str:
  23. """
  24. Note: This only demangles names that were mangled by this specific
  25. PackageMangler. It will pass through names created by a different
  26. PackageMangler instance.
  27. """
  28. if mangled.startswith(self._mangle_parent + "."):
  29. return mangled.partition(".")[2]
  30. # wasn't a mangled name
  31. return mangled
  32. def parent_name(self):
  33. return self._mangle_parent
  34. def is_mangled(name: str) -> bool:
  35. return bool(re.match(r"<torch_package_\d+>", name))
  36. def demangle(name: str) -> str:
  37. """
  38. Note: Unlike PackageMangler.demangle, this version works on any
  39. mangled name, irrespective of which PackageMangler created it.
  40. """
  41. if is_mangled(name):
  42. first, sep, last = name.partition(".")
  43. # If there is only a base mangle prefix, e.g. '<torch_package_0>',
  44. # then return an empty string.
  45. return last if len(sep) != 0 else ""
  46. return name
  47. def get_mangle_prefix(name: str) -> str:
  48. return name.partition(".")[0] if is_mangled(name) else name