relationalai 0.13.5__py3-none-any.whl → 1.0.0a2__py3-none-any.whl

relationalai_test_util/snapshot.py

@@ -1,252 +0,0 @@
-from __future__ import annotations
-import io
-import logging
-import random
-import os
-from pathlib import Path
-import textwrap
-from typing import Any, Callable, TypedDict
-
-from pandas import DataFrame
-import pytest
-
-import rich
-from rich.syntax import Syntax
-
-import relationalai as rai
-from relationalai.clients import config as cfg
-from relationalai import debugging
-from relationalai.debugging import logger
-from relationalai.errors import RAIException
-from relationalai.metamodel import Task
-
-################################################################################
-# Parameterize tests over all files in a directory
-################################################################################
-
-def path_to_slug(path: Path, base_path:str|Path):
-    return str(path.relative_to(base_path)).replace("/", "__").replace(".py", "")
-
-def all_files_in_dir(test_case_dir: Path, *, pattern = "*.py"):
-    "Decorator to parameterize a pytest over files in the given directory."
-    test_case_files = [path for path in test_case_dir.rglob(pattern)]
-    return pytest.mark.parametrize(
-        "file_path", test_case_files, ids=lambda path: path_to_slug(path, test_case_dir)
-    )
-
-################################################################################
-# Tap block info out of the debug log
-################################################################################
-
-class ISourceInfo(TypedDict):
-    file: str
-    line: int
-    block: str
-
-class IPassInfo(TypedDict):
-    name: str
-    task: str
-    elapsed: float
-
-class ISampledResults(TypedDict):
-    values: DataFrame
-    count: int
-
-class Block(TypedDict):
-    task: Task
-    source: ISourceInfo
-    passes: list[IPassInfo]
-    emitted: str
-    emit_time: float
-    results: ISampledResults|None
-    alt_format_results: Any|None
-    elapsed: float|None
-
-class QueryTextLogger(logging.Handler):
-    def __init__(self):
-        super().__init__()
-        # compilation events that get 'results' added to them
-        self.blocks: list[Block] = []
-
-    def emit(self, record):
-        d = record.msg
-        if isinstance(d, dict):
-            if d["event"] == "compilation":
-                self.blocks.append({
-                    "task": d["task"],
-                    "source": d["source"],
-                    "passes": d["passes"],
-                    "emitted": d["emitted"],
-                    "emit_time": d["emit_time"],
-                    "results": None,
-                    "alt_format_results": None,
-                    "elapsed": None
-                })
-            elif d["event"] == "time" and d["type"] == "query" and not d.get("internal"):
-                last = self.blocks[-1]
-                last["results"] = d["results"]
-                last["alt_format_results"] = d.get("alt_format_results")
-                last["elapsed"] = d["elapsed"]
-
-################################################################################
-# Snapshot validation
-################################################################################
-
-class SnapshotError(Exception):
-    def __init__(self, file_path: Path, code: str, task: Task, header: str|None = None, footer: str|None = None):
-        self.file_path = file_path
-        self.source_code = code
-        self.task = task
-        self.header = header
-        self.footer = footer
-
-    def __str__(self):
-        file_path = self.file_path
-
-        with io.StringIO() as buf:
-            console = rich.console.Console(file=buf, force_terminal=True)
-            if self.header:
-                console.print(self.header)
-
-            source_info = debugging.get_source(self.task)
-            assert source_info and source_info.line is not None
-            source = debugging.find_block_in(self.source_code, source_info.line, str(file_path))
-
-            base = os.getcwd()
-            console.print("\nIn", f"./{file_path.relative_to(base)}" if file_path.is_relative_to(base) else file_path)
-            if source.source:
-                console.print(Syntax(source.source, "python", dedent=True, line_numbers=True, start_line=source.line, padding=1))
-
-            if self.footer:
-                console.print("\n" + self.footer)
-            v = buf.getvalue()
-        return v
-
-    def __repr__(self):
-        return self.__str__()
-
-class ChangedSnapshotError(SnapshotError):
-    def __init__(self, file_path: Path, code: str, task: Task, snap_file: str, diff: str):
-        header = f"The value for the following block no longer matches its snapshot ({snap_file}):"
-        footer = f"[bold]Diff[/bold]\n{diff}\n\n" + \
-        "[yellow]If this change is expected, you can update the snapshot by running pytest with [bold]--snapshot-update[/bold][/yellow]"
-        super().__init__(file_path, code, task, header, footer)
-
-class NewSnapshotError(SnapshotError):
-    def __init__(self, file_path: Path, code: str, task: Task, value: str|None):
-        header = "No snapshot exists for the following block:"
-        footer = textwrap.dedent(f"""
-        Got:
-        {value}
-        [yellow]To save the current result as a snapshot, run pytest with [bold]--snapshot-update[/bold][/yellow]
-        """)
-        super().__init__(file_path, code, task, header, footer)
-
-
-def exec_traced(code: str, model_kwargs:dict|None = None, file_path: str|None = None) -> list[Block]:
-    # install logger to capture Rel compilation
-    query_logger = QueryTextLogger()
-    logger.addHandler(query_logger)
-
-
-    model_kwargs = model_kwargs or {}
-    if file_path:
-        # ensure that the exec knows what file is being run
-        model_kwargs["__file__"] = str(file_path)
-
-    try:
-        if file_path:
-            code_object = compile(code, file_path, 'exec')
-            exec(code_object, model_kwargs)
-        else:
-            exec(code, model_kwargs)
-    except RAIException as err:
-        err.pprint() # @FIXME: Is this still needed?
-        raise err
-    finally:
-        logger.removeHandler(query_logger)
-
-    return query_logger.blocks
-
-def validate_block_snapshots(
-        file_path: Path,
-        snapshot,
-        get_snapshot_str: Callable, # given a block, return the snapshot string
-        snapshot_prefix:str,
-        model_kwargs:dict|None = None
-    ):
-    """Execute the program at `file_path`, applying `get_snapshot_str` to all
-    blocks and ensuring the result has not changed since the last run."""
-    with open(file_path, "r") as file:
-        code = file.read()
-        # @TODO: Consider suppressing stdout
-
-        blocks = exec_traced(code, model_kwargs, str(file_path))
-        block_index = 0
-        for block in blocks:
-            snapshot_path = f"{snapshot_prefix}{block_index}.txt"
-            cmp_string = None
-            try:
-                cmp_string = get_snapshot_str(block)
-                if cmp_string is not None:
-                    snapshot.assert_match(cmp_string, snapshot_path)
-                    # only counting blocks that actually assert something
-                    block_index += 1
-            except RAIException as err:
-                err.pprint()
-                raise err from None
-            except AssertionError as err:
-                lines = str(err).splitlines()
-                if "doesn't exist" in lines[0]:
-                    raise NewSnapshotError(file_path, code, block["task"], cmp_string) from None
-
-                if len(lines) >= 3 and "does not match" in lines[0]:
-                    diff = '\n'.join(lines[3:])
-                    raise ChangedSnapshotError(file_path, code, block["task"], snapshot_path, diff) from None
-
-                raise
-
-
-def randomize(name: str) -> str:
-    return f"{name}_{random.randint(1000000000, 9999999999)}"
-
-def do_snapshot_test(
-        snapshot,
-        file_path: Path,
-        config: cfg.Config|None = None,
-        engine_size: str|None = None,
-        runs: int = 1,
-        outer_span_type: str = "test",
-        create_db: bool = True,
-    ):
-    """Execute the program at `file_path`, validating that query results
-    have not changed since the last run. Uses the given `config` and
-    `engine_size` if provided, or the ambient config if not."""
-    test_name = file_path.stem
-
-    provider = rai.Resources(config=config)
-    config = provider.config
-
-    def get_results(event: Block):
-        if event["alt_format_results"]:
-            return str(event["alt_format_results"].to_pandas())
-        if event["results"]:
-            return str(event["results"]["values"])
-
-    with debugging.span(outer_span_type, name=test_name, engine_size=engine_size):
-        for i in range(runs):
-            with debugging.span("run", idx=i):
-                program_span = debugging.get_program_span()
-                db_name = randomize(test_name)
-                validate_block_snapshots(file_path, snapshot, get_results, "query", {
-                    "name": db_name,
-                    "config": config,
-                    "span": program_span,
-                })
-                # Some tests always use the same DB, never creating a DB named `db_name`,
-                # so we don't need to delete it.
-                if create_db:
-                    try:
-                        provider.delete_graph(db_name)
-                    except Exception as e:
-                        print(f"Failed to delete graph {db_name}: {e}")

relationalai_test_util/traceback.py

@@ -1,118 +0,0 @@
-import os
-import sys
-import re
-from pathlib import Path
-import traceback
-from typing import Union
-from colorama import Style
-
-################################################################################
-# Relative path selection
-################################################################################
-
-
-def find_project_root(path):
-    path = Path(path).resolve()
-
-    # Common project root indicators
-    root_indicators = [
-        "pyproject.toml",
-        "setup.py",
-        "setup.cfg",
-        ".git",
-        ".hg",
-        ".svn",
-        "requirements.txt",
-        "Pipfile",
-        "poetry.lock",
-    ]
-
-    current_path = path
-    while current_path != current_path.parent:
-        if any((current_path / indicator).exists() for indicator in root_indicators):
-            return current_path
-        current_path = current_path.parent
-
-    return None
-
-
-PROJECT_ROOT = find_project_root(os.getcwd())
-
-def relative_to_project(path: Path):
-    if PROJECT_ROOT and path.is_relative_to(PROJECT_ROOT):
-        return path.relative_to(PROJECT_ROOT)
-
-
-def relative_to_venv(path: Path):
-    if hasattr(sys, "prefix"):
-        venv_path = Path(sys.prefix)
-        if path.is_relative_to(venv_path):
-            return path.relative_to(venv_path.parent)
-
-
-def relative_to_site_packages(path: Path):
-    best: Union[Path, None] = None
-
-    for site_path in sys.path:
-        if "site-packages" in site_path:
-            site_path = Path(site_path)
-            if path.is_relative_to(site_path):
-                new = path.relative_to(site_path.parent)
-                if not best or len(str(best)) > len(str(new)):
-                    best = new
-
-    return best
-
-
-def condense_path(path):
-    resolved = Path(path).resolve()
-    best: Union[Path, None] = None
-
-    for try_option in (
-        relative_to_project,
-        relative_to_venv,
-        relative_to_site_packages,
-    ):
-        option = try_option(resolved)
-        if not best:
-            best = option
-        elif option and len(str(option)) < len(str(best)):
-            best = option
-
-    if best:
-        return str(best)
-    else:
-        return str(path)
-
-
-################################################################################
-# Pretty traceback
-################################################################################
-
-def condense_traceback(exception: Exception, filter_frame_patterns = ["site-packages/_pytest", "site-packages/pluggy"]):
-    tb = traceback.extract_tb(exception.__traceback__)
-    max_prefix_size = 0
-    prefixes: list[str] = []
-    linenos: list[int] = []
-    messages: list[str] = []
-    for filename, lineno, name, text in tb:
-        prefix = f"{condense_path(filename)} in {name}"
-        exclude = False
-        for pattern in filter_frame_patterns:
-            if re.search(pattern, prefix):
-                exclude = True
-        if exclude:
-            continue
-
-        if len(prefix) > max_prefix_size:
-            max_prefix_size = len(prefix)
-        prefixes.append(prefix)
-        linenos.append(lineno)
-        messages.append(text.strip())
-
-    return "\n".join(
-        [
-            f"{Style.DIM}{prefix: <{max_prefix_size}} {lineno: >4} |{Style.RESET_ALL} {msg}"
-            for prefix, lineno, msg in zip(prefixes, linenos, messages)
-        ]
-    )