132 lines
3.8 KiB
Python
132 lines
3.8 KiB
Python
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)
|