Search modules and packages in arcadia source tree. See https://wiki.yandex-team.ru/devtools/extended-python-source-search/ for details
| 536 | |
| 537 | |
| 538 | class ArcadiaSourceFinder: |
| 539 | """ |
| 540 | Search modules and packages in arcadia source tree. |
| 541 | See https://wiki.yandex-team.ru/devtools/extended-python-source-search/ for details |
| 542 | """ |
| 543 | NAMESPACE_PREFIX = b'py/namespace/' |
| 544 | PY_EXT = '.py' |
| 545 | YA_MAKE = 'ya.make' |
| 546 | S_IFDIR = 0o040000 |
| 547 | |
| 548 | def __init__(self, source_root): |
| 549 | self.source_root = source_root |
| 550 | self.module_path_cache = {'': set()} |
| 551 | for key, dirty_path in iter_keys(self.NAMESPACE_PREFIX): |
| 552 | # dirty_path contains unique prefix to prevent repeatable keys in the resource storage |
| 553 | path = dirty_path.split(b'/', 1)[1] |
| 554 | namespaces = find(key).split(b':') |
| 555 | for n in namespaces: |
| 556 | package_name = _s(n.rstrip(b'.')) |
| 557 | self.module_path_cache.setdefault(package_name, set()).add(_s(path)) |
| 558 | # Fill parents with default empty path set if parent doesn't exist in the cache yet |
| 559 | while package_name: |
| 560 | package_name = package_name.rpartition('.')[0] |
| 561 | if package_name in self.module_path_cache: |
| 562 | break |
| 563 | self.module_path_cache.setdefault(package_name, set()) |
| 564 | for package_name in self.module_path_cache.keys(): |
| 565 | self._add_parent_dirs(package_name, visited=set()) |
| 566 | |
| 567 | def get_module_path(self, fullname): |
| 568 | """ |
| 569 | Find file path for module 'fullname'. |
| 570 | For packages caller pass fullname as 'package.__init__'. |
| 571 | Return None if nothing is found. |
| 572 | """ |
| 573 | try: |
| 574 | if not self.is_package(fullname): |
| 575 | return _b(self._cache_module_path(fullname)) |
| 576 | except ImportError: |
| 577 | pass |
| 578 | |
| 579 | def is_package(self, fullname): |
| 580 | """Check if fullname is a package. Raise ImportError if fullname is not found""" |
| 581 | path = self._cache_module_path(fullname) |
| 582 | if isinstance(path, set): |
| 583 | return True |
| 584 | if isinstance(path, str): |
| 585 | return False |
| 586 | raise ImportError(fullname) |
| 587 | |
| 588 | def iter_modules(self, package_prefix, prefix): |
| 589 | paths = self._cache_module_path(package_prefix.rstrip('.')) |
| 590 | if paths is not None: |
| 591 | # Note: it's ok to yield duplicates because pkgutil discards them |
| 592 | |
| 593 | # Yield from cache |
| 594 | import re |
| 595 | rx = re.compile(re.escape(package_prefix) + r'([^.]+)$') |