glob_group.py 3.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. import re
  2. from typing import Iterable, Union
  3. GlobPattern = Union[str, Iterable[str]]
  4. class GlobGroup:
  5. """A set of patterns that candidate strings will be matched against.
  6. A candidate is composed of a list of segments separated by ``separator``, e.g. "foo.bar.baz".
  7. A pattern contains one or more segments. Segments can be:
  8. - A literal string (e.g. "foo"), which matches exactly.
  9. - A string containing a wildcard (e.g. "torch*", or "foo*baz*"). The wildcard matches
  10. any string, including the empty string.
  11. - A double wildcard ("**"). This matches against zero or more complete segments.
  12. Examples:
  13. ``torch.**``: matches ``torch`` and all its submodules, e.g. ``torch.nn`` and ``torch.nn.functional``.
  14. ``torch.*``: matches ``torch.nn`` or ``torch.functional``, but not ``torch.nn.functional``.
  15. ``torch*.**``: matches ``torch``, ``torchvision``, and all their submodules.
  16. A candidates will match the ``GlobGroup`` if it matches any of the ``include`` patterns and
  17. none of the ``exclude`` patterns.
  18. Args:
  19. include (Union[str, Iterable[str]]): A string or list of strings,
  20. each representing a pattern to be matched against. A candidate
  21. will match if it matches *any* include pattern
  22. exclude (Union[str, Iterable[str]]): A string or list of strings,
  23. each representing a pattern to be matched against. A candidate
  24. will be excluded from matching if it matches *any* exclude pattern.
  25. separator (str): A string that delimits segments in candidates and
  26. patterns. By default this is "." which corresponds to how modules are
  27. named in Python. Another common value for this is "/", which is
  28. the Unix path separator.
  29. """
  30. def __init__(
  31. self, include: GlobPattern, *, exclude: GlobPattern = (), separator: str = "."
  32. ):
  33. self._dbg = f"GlobGroup(include={include}, exclude={exclude})"
  34. self.include = GlobGroup._glob_list(include, separator)
  35. self.exclude = GlobGroup._glob_list(exclude, separator)
  36. self.separator = separator
  37. def __str__(self):
  38. return self._dbg
  39. def __repr__(self):
  40. return self._dbg
  41. def matches(self, candidate: str) -> bool:
  42. candidate = self.separator + candidate
  43. return any(p.fullmatch(candidate) for p in self.include) and all(
  44. not p.fullmatch(candidate) for p in self.exclude
  45. )
  46. @staticmethod
  47. def _glob_list(elems: GlobPattern, separator: str = "."):
  48. if isinstance(elems, str):
  49. return [GlobGroup._glob_to_re(elems, separator)]
  50. else:
  51. return [GlobGroup._glob_to_re(e, separator) for e in elems]
  52. @staticmethod
  53. def _glob_to_re(pattern: str, separator: str = "."):
  54. # to avoid corner cases for the first component, we prefix the candidate string
  55. # with '.' so `import torch` will regex against `.torch`, assuming '.' is the separator
  56. def component_to_re(component):
  57. if "**" in component:
  58. if component == "**":
  59. return "(" + re.escape(separator) + "[^" + separator + "]+)*"
  60. else:
  61. raise ValueError("** can only appear as an entire path segment")
  62. else:
  63. return re.escape(separator) + ("[^" + separator + "]*").join(
  64. re.escape(x) for x in component.split("*")
  65. )
  66. result = "".join(component_to_re(c) for c in pattern.split(separator))
  67. return re.compile(result)