dbt-selly/dbt-env/lib/python3.8/site-packages/dbt/parser/search.py

import os
from dataclasses import dataclass
from typing import (
    List, Callable, Iterable, Set, Union, Iterator, TypeVar, Generic
)

from dbt.clients.jinja import extract_toplevel_blocks, BlockTag
from dbt.clients.system import find_matching
from dbt.config import Project
from dbt.contracts.files import FilePath, AnySourceFile
from dbt.exceptions import CompilationException, InternalException


# What's the point of wrapping a SourceFile with this class?
# Could it be removed?
@dataclass
class FileBlock:
    file: AnySourceFile

    @property
    def name(self):
        base = os.path.basename(self.file.path.relative_path)
        name, _ = os.path.splitext(base)
        return name

    @property
    def contents(self):
        return self.file.contents

    @property
    def path(self):
        return self.file.path


# The BlockTag is used in Jinja processing
# Why do we have different classes where the only
# difference is what 'contents' returns?
@dataclass
class BlockContents(FileBlock):
    file: AnySourceFile  # if you remove this, mypy will get upset
    block: BlockTag

    @property
    def name(self):
        return self.block.block_name

    @property
    def contents(self):
        return self.block.contents


@dataclass
class FullBlock(FileBlock):
    file: AnySourceFile  # if you remove this, mypy will get upset
    block: BlockTag

    @property
    def name(self):
        return self.block.block_name

    @property
    def contents(self):
        return self.block.full_block


class FilesystemSearcher(Iterable[FilePath]):
    def __init__(
        self, project: Project, relative_dirs: List[str], extension: str
    ) -> None:
        self.project = project
        self.relative_dirs = relative_dirs
        self.extension = extension

    def __iter__(self) -> Iterator[FilePath]:
        ext = "[!.#~]*" + self.extension

        root = self.project.project_root

        for result in find_matching(root, self.relative_dirs, ext):
            if 'searched_path' not in result or 'relative_path' not in result:
                raise InternalException(
                    'Invalid result from find_matching: {}'.format(result)
                )
            file_match = FilePath(
                searched_path=result['searched_path'],
                relative_path=result['relative_path'],
                modification_time=result['modification_time'],
                project_root=root,
            )
            yield file_match


Block = Union[BlockContents, FullBlock]

BlockSearchResult = TypeVar('BlockSearchResult', BlockContents, FullBlock)

BlockSearchResultFactory = Callable[[AnySourceFile, BlockTag], BlockSearchResult]


class BlockSearcher(Generic[BlockSearchResult], Iterable[BlockSearchResult]):
    def __init__(
        self,
        source: List[FileBlock],
        allowed_blocks: Set[str],
        source_tag_factory: BlockSearchResultFactory
    ) -> None:
        self.source = source
        self.allowed_blocks = allowed_blocks
        self.source_tag_factory: BlockSearchResultFactory = source_tag_factory

    def extract_blocks(self, source_file: FileBlock) -> Iterable[BlockTag]:
        try:
            blocks = extract_toplevel_blocks(
                source_file.contents,
                allowed_blocks=self.allowed_blocks,
                collect_raw_data=False
            )
            # this makes mypy happy, and this is an invariant we really need
            for block in blocks:
                assert isinstance(block, BlockTag)
                yield block

        except CompilationException as exc:
            if exc.node is None:
                exc.add_node(source_file)
            raise

    def __iter__(self) -> Iterator[BlockSearchResult]:
        for entry in self.source:
            for block in self.extract_blocks(entry):
                yield self.source_tag_factory(entry.file, block)