ocalc 0.1.0__py3-none-any.whl

ocalc/__init__.py

@@ -0,0 +1,216 @@
+"""Measure a Python function and estimate its time and space complexity."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from ocalc.bench import BenchResult, Point, benchmark
+from ocalc.comment import (
+    Protocol,
+    apply_comment,
+    existing_comment,
+    format_comment,
+    format_table,
+)
+from ocalc.fit import FitResult, fit_curve
+from ocalc.generate import (
+    FunctionNotFound,
+    NeedSample,
+    callable_location,
+    load_function,
+    parse_sample,
+    plan_inputs,
+)
+from ocalc.static import FunctionEstimate, SourceAnalysis, analyze_source
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "FunctionEstimate",
+    "FunctionNotFound",
+    "NeedSample",
+    "Report",
+    "SourceAnalysis",
+    "analyze_source",
+    "annotate_file",
+    "measure",
+    "parse_sample",
+]
+
+
+@dataclass(frozen=True)
+class Report:
+    comment: str
+    points: list[Point]
+    time: FitResult
+    space: FitResult
+    worst_s: float
+    median_s: float
+    worst_n: int | None
+    protocol: Protocol
+
+
+def measure(
+    func,
+    *sample_args,
+    budget: float | None = None,
+    min_n: int | None = None,
+    max_n: int | None = None,
+    measures: int | None = None,
+    repeats: int | None = None,
+    file: str | None = None,
+    qualname: str | None = None,
+    protocol: Protocol | None = None,
+    previous_s: float | None = None,
+    **sample_kwargs,
+) -> Report:
+    """Measure a function and return its complexity report.
+
+    Args:
+        func: The callable to measure.
+        sample_args: Example positional arguments. The first collection grows with ``n``.
+        budget: Seconds to spend measuring. The default is 5.
+        min_n: Smallest input size.
+        max_n: Largest input size.
+        measures: How many sizes to try.
+        repeats: Timed calls at each size.
+        file: Source file, when the callable should be re-imported for the timed run.
+        qualname: Qualname used to re-import ``file``.
+        protocol: Stored sizes and budget from an earlier comment.
+        previous_s: Worst time from that earlier comment.
+        sample_kwargs: Example keyword arguments.
+
+    Returns:
+        The comment, the size table, and the fitted time and space classes.
+    """
+    plan = plan_inputs(func, sample_args, sample_kwargs)
+    if protocol is not None and protocol.constant and plan.mode == "constant":
+        if budget is None:
+            budget = protocol.budget
+        if repeats is None:
+            repeats = protocol.repeats
+    elif protocol is not None and not protocol.constant and plan.mode != "constant":
+        if budget is None:
+            budget = protocol.budget
+        if min_n is None:
+            min_n = protocol.min_n
+        if max_n is None:
+            max_n = protocol.max_n
+        if measures is None:
+            measures = protocol.measures
+        if repeats is None:
+            repeats = protocol.repeats
+
+    located_file, located_name = callable_location(func)
+    file = file or located_file
+    qualname = qualname or located_name
+    result = benchmark(
+        func,
+        plan,
+        budget=5.0 if budget is None else budget,
+        min_n=min_n,
+        max_n=max_n,
+        measures=measures,
+        repeats=repeats,
+        file=file,
+        qualname=qualname,
+    )
+    time_fit, space_fit = _fits(result)
+    comment = format_comment(
+        time=time_fit,
+        space=space_fit,
+        worst_s=result.worst_s,
+        worst_n=result.worst_n,
+        median_s=result.median_s,
+        protocol=result.protocol,
+        previous_s=previous_s,
+    )
+    return Report(
+        comment=comment,
+        points=result.points,
+        time=time_fit,
+        space=space_fit,
+        worst_s=result.worst_s,
+        median_s=result.median_s,
+        worst_n=result.worst_n,
+        protocol=result.protocol,
+    )
+
+
+def annotate_file(
+    path,
+    qualname: str,
+    *sample_args,
+    budget: float | None = None,
+    min_n: int | None = None,
+    max_n: int | None = None,
+    measures: int | None = None,
+    repeats: int | None = None,
+    **sample_kwargs,
+) -> Report:
+    """Benchmark a function in a file and write the comment above it.
+
+    A later call reuses the protocol stored in the existing comment and records
+    the previous worst time.
+
+    Args:
+        path: Python file that defines the function.
+        qualname: Function qualname, such as ``two_sum`` or ``Box.run``.
+        sample_args: Example positional arguments.
+        budget: Seconds to spend measuring.
+        min_n: Smallest input size.
+        max_n: Largest input size.
+        measures: How many sizes to try.
+        repeats: Timed calls at each size.
+        sample_kwargs: Example keyword arguments.
+
+    Returns:
+        The report that was written into the file.
+    """
+    file_path = Path(path)
+    source = file_path.read_text(encoding="utf-8")
+    previous = existing_comment(source, qualname)
+    protocol = previous.protocol if previous is not None else None
+    previous_s = previous.worst_s if previous is not None else None
+    func = load_function(file_path, qualname)
+    report = measure(
+        func,
+        *sample_args,
+        budget=budget,
+        min_n=min_n,
+        max_n=max_n,
+        measures=measures,
+        repeats=repeats,
+        file=str(file_path.resolve()),
+        qualname=qualname,
+        protocol=protocol,
+        previous_s=previous_s,
+        **sample_kwargs,
+    )
+    updated = apply_comment(source, qualname, report.comment)
+    with file_path.open("w", encoding="utf-8", newline="") as handle:
+        handle.write(updated)
+    return report
+
+
+def _fits(result: BenchResult) -> tuple[FitResult, FitResult]:
+    if result.constant:
+        done = FitResult("1", "O(1)", "high")
+        return done, done
+    sizes = [point.n for point in result.points]
+    times = [point.time_s for point in result.points]
+    time_fit = fit_curve(sizes, times, exponential=result.exponential)
+    return time_fit, _space_fit(sizes, [point.memory_bytes for point in result.points])
+
+
+def _space_fit(sizes, memory) -> FitResult:
+    used = np.asarray(memory, dtype=float)
+    if used.size == 0 or np.nanmax(used) <= 0:
+        return FitResult("1", "O(1)", "high")
+    positive = used[used > 0]
+    floor = float(positive.min()) if positive.size else 1.0
+    used = np.where(used > 0, used, floor)
+    return fit_curve(sizes, used)

ocalc/bench.py

@@ -0,0 +1,287 @@
+"""Time a function across a size grid and record peak Python allocations."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import math
+import multiprocessing
+import time
+import tracemalloc
+from dataclasses import dataclass
+from statistics import median
+
+from ocalc.comment import Protocol
+from ocalc.generate import InputPlan, load_function
+
+DEFAULT_BUDGET = 5.0
+DEFAULT_MEASURES = 8
+DEFAULT_MIN_N = 100
+DEFAULT_MAX_N = 100_000
+DEFAULT_MIN_N_INT = 8
+_REPEATS = 5
+
+
+@dataclass(frozen=True)
+class Point:
+    n: int | None
+    time_s: float
+    memory_bytes: float
+
+
+@dataclass
+class BenchResult:
+    points: list[Point]
+    protocol: Protocol
+    exponential: bool
+    constant: bool
+    timed_out: bool
+    worst_s: float
+    median_s: float
+    worst_n: int | None
+
+
+def geometric_sizes(min_n: int, max_n: int, measures: int) -> list[int]:
+    """Evenly spaced sizes on a log scale, including both ends."""
+    if measures < 1:
+        raise ValueError("measures must be at least 1")
+    if min_n > max_n:
+        raise ValueError("min_n cannot be greater than max_n")
+    if measures == 1 or min_n == max_n:
+        return [int(max_n)]
+    ratio = (max_n / min_n) ** (1 / (measures - 1))
+    raw = [int(round(min_n * ratio**i)) for i in range(measures)]
+    raw[0] = int(min_n)
+    raw[-1] = int(max_n)
+    sizes: list[int] = []
+    for size in raw:
+        size = max(int(min_n), size)
+        if not sizes or size > sizes[-1]:
+            sizes.append(size)
+    if sizes[-1] != int(max_n):
+        sizes.append(int(max_n))
+    return sizes
+
+
+def benchmark(
+    func,
+    plan: InputPlan,
+    *,
+    budget: float,
+    min_n: int | None,
+    max_n: int | None,
+    measures: int | None,
+    repeats: int | None,
+    file: str | None,
+    qualname: str | None,
+) -> BenchResult:
+    measures = DEFAULT_MEASURES if measures is None else measures
+    budget = DEFAULT_BUDGET if budget is None else budget
+    repeats = _REPEATS if repeats is None else max(1, int(repeats))
+    if plan.mode == "constant":
+        return _benchmark_constant(func, plan, budget, repeats, file, qualname)
+
+    if plan.mode == "int_n":
+        min_n = DEFAULT_MIN_N_INT if min_n is None else min_n
+        max_n = DEFAULT_MAX_N if max_n is None else max_n
+    else:
+        min_n = DEFAULT_MIN_N if min_n is None else min_n
+        max_n = DEFAULT_MAX_N if max_n is None else max_n
+    protocol = Protocol(int(min_n), int(max_n), int(measures), float(budget), repeats=repeats)
+    sizes = geometric_sizes(protocol.min_n, protocol.max_n, protocol.measures)
+    points: list[Point] = []
+    exponential = False
+    timed_out = False
+    started = time.perf_counter()
+
+    for size in sizes:
+        remaining = budget - (time.perf_counter() - started)
+        if remaining <= 0.05:
+            break
+        if len(points) >= 2 and _explosive(
+            points[-2].n, points[-2].time_s, points[-1].n, points[-1].time_s
+        ):
+            exponential = True
+            break
+        if points:
+            ratio = size / points[-1].n
+            predicted = points[-1].time_s * max(ratio, 1.0) ** 2
+            if predicted > remaining:
+                break
+        args, kwargs = plan.build(size)
+        status, times, peak = _measure_call(
+            func,
+            args,
+            kwargs,
+            file=file,
+            qualname=qualname,
+            repeats=repeats,
+            timeout=max(0.2, remaining),
+        )
+        if status == "timeout":
+            timed_out = True
+            if points and (
+                points[-1].n <= 40
+                and points[-1].time_s > 1e-4
+                or (
+                    len(points) >= 2
+                    and _explosive(points[-2].n, points[-2].time_s, points[-1].n, points[-1].time_s)
+                )
+            ):
+                exponential = True
+            break
+        if status == "error":
+            raise RuntimeError(times)
+        point = Point(size, float(median(times)), float(peak))
+        points.append(point)
+        if len(points) >= 2 and _explosive(
+            points[-2].n, points[-2].time_s, points[-1].n, points[-1].time_s
+        ):
+            exponential = True
+            break
+
+    if not points:
+        raise RuntimeError(
+            "No measurement finished within the budget. Raise --budget or lower --min-n."
+        )
+    slowest = max(points, key=lambda item: item.time_s)
+    return BenchResult(
+        points=points,
+        protocol=protocol,
+        exponential=exponential,
+        constant=False,
+        timed_out=timed_out,
+        worst_s=slowest.time_s,
+        median_s=float(median(item.time_s for item in points)),
+        worst_n=slowest.n,
+    )
+
+
+def _benchmark_constant(func, plan, budget, repeats, file, qualname) -> BenchResult:
+    args, kwargs = plan.build(0)
+    status, times, peak = _measure_call(
+        func,
+        args,
+        kwargs,
+        file=file,
+        qualname=qualname,
+        repeats=repeats,
+        timeout=max(0.2, budget),
+    )
+    if status == "timeout":
+        raise RuntimeError(
+            "No measurement finished within the budget. Raise --budget or lower --min-n."
+        )
+    if status == "error":
+        raise RuntimeError(times)
+    protocol = Protocol(0, 0, 0, float(budget), constant=True, repeats=repeats)
+    point = Point(None, float(median(times)), float(peak))
+    return BenchResult(
+        points=[point],
+        protocol=protocol,
+        exponential=False,
+        constant=True,
+        timed_out=False,
+        worst_s=float(max(times)),
+        median_s=float(median(times)),
+        worst_n=None,
+    )
+
+
+def _explosive(n0, t0, n1, t1) -> bool:
+    if t0 <= 0 or t1 <= 0 or n1 <= n0:
+        return False
+    n_ratio = n1 / n0
+    t_ratio = t1 / t0
+    if n_ratio < 1.5 and t_ratio > 2 and n1 <= 64 and t1 > 1e-4:
+        return True
+    if t1 < 1e-3 or n_ratio <= 1:
+        return False
+    exponent = math.log(t_ratio) / math.log(n_ratio)
+    return exponent > 5
+
+
+def _call(func, args, kwargs):
+    result = func(*args, **kwargs)
+    if inspect.iscoroutine(result):
+        return asyncio.run(result)
+    return result
+
+
+def _measure_call(func, args, kwargs, *, file, qualname, repeats, timeout):
+    if file and qualname:
+        return _run_process(file, qualname, args, kwargs, repeats, timeout)
+    times, peak = collect_sample(func, args, kwargs, repeats)
+    return "ok", times, peak
+
+
+def collect_sample(func, args, kwargs, repeats: int = _REPEATS):
+    """Warm up, time several calls, then record one tracemalloc peak."""
+    _call(func, args, kwargs)
+    batch = 1
+
+    def once() -> float:
+        started = time.perf_counter()
+        for _ in range(batch):
+            _call(func, args, kwargs)
+        return (time.perf_counter() - started) / batch
+
+    first = once()
+    if first < 0.001:
+        batch = min(10_000, max(1, math.ceil(0.002 / max(first, 1e-9))))
+        first = once()
+    times = [first]
+    if first <= 0.2:
+        for _ in range(repeats - 1):
+            duration = once()
+            times.append(duration)
+            if duration > 0.2:
+                break
+    tracemalloc.start()
+    try:
+        _call(func, args, kwargs)
+        _current, peak = tracemalloc.get_traced_memory()
+    finally:
+        tracemalloc.stop()
+    return times, float(peak)
+
+
+def _run_process(file, qualname, args, kwargs, repeats, timeout):
+    context = multiprocessing.get_context("spawn")
+    parent, child = context.Pipe(duplex=False)
+    process = context.Process(
+        target=_worker,
+        args=(child, str(file), qualname, args, kwargs, repeats),
+        daemon=True,
+    )
+    process.start()
+    child.close()
+    try:
+        process.join(timeout)
+        if process.is_alive():
+            process.terminate()
+            process.join()
+            return "timeout", None, None
+        if not parent.poll(0.2):
+            return "error", "The measurement process stopped without a result.", None
+        message = parent.recv()
+    finally:
+        parent.close()
+    if not message or message[0] != "ok":
+        detail = message[1] if message and len(message) > 1 else "measurement failed"
+        return "error", detail, None
+    return "ok", message[1], message[2]
+
+
+def _worker(conn, file, qualname, args, kwargs, repeats):
+    try:
+        func = load_function(file, qualname)
+        times, peak = collect_sample(func, args, kwargs, repeats)
+        conn.send(("ok", times, peak))
+    except Exception as exc:
+        try:
+            conn.send(("err", f"{type(exc).__name__}: {exc}"))
+        except Exception:
+            pass
+    finally:
+        conn.close()

ocalc/cli.py

@@ -0,0 +1,86 @@
+"""Command line for measuring and annotating functions."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+
+from ocalc import FunctionNotFound, NeedSample, annotate_file, measure, parse_sample
+from ocalc.comment import format_metrics, format_table
+from ocalc.generate import load_function
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="ocalc",
+        description="Measure a Python function and estimate its time and space complexity.",
+    )
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    measure_cmd = commands.add_parser("measure", help="Print a complexity report.")
+    annotate_cmd = commands.add_parser("annotate", help="Write the complexity comment above a function.")
+    for command in (measure_cmd, annotate_cmd):
+        command.add_argument("file", help="Python file that defines the function.")
+        command.add_argument("-f", "--function", required=True, help="Function qualname, such as two_sum.")
+        command.add_argument(
+            "--sample",
+            default="",
+            help='Literal arguments, for example "[1, 2, 3], 9".',
+        )
+        protocol = command.add_argument_group("protocol")
+        protocol.add_argument("--budget", type=float, default=None, help="Seconds to spend measuring. Default: 5.")
+        protocol.add_argument("--min-n", type=int, default=None, help="Smallest input size. Default: 100, or 8 for a single int.")
+        protocol.add_argument("--max-n", type=int, default=None, help="Largest input size. Default: 100000.")
+        protocol.add_argument("--measures", type=int, default=None, help="How many sizes to try. Default: 8.")
+        protocol.add_argument("--repeats", type=int, default=None, help="Timed calls at each size. Default: 5.")
+    commands.add_parser("mcp", help="Serve the measure and annotate tools over stdio.")
+
+    args = parser.parse_args(argv)
+    if args.command == "mcp":
+        from ocalc.mcp_server import serve
+
+        serve()
+        return 0
+    try:
+        return _run(args)
+    except NeedSample as exc:
+        print(exc, file=sys.stderr)
+        return 2
+    except FunctionNotFound as exc:
+        print(exc, file=sys.stderr)
+        return 2
+    except (RuntimeError, OSError) as exc:
+        print(exc, file=sys.stderr)
+        return 1
+
+
+def _run(args) -> int:
+    sample_args, sample_kwargs = parse_sample(args.sample)
+    options = dict(
+        budget=args.budget,
+        min_n=args.min_n,
+        max_n=args.max_n,
+        measures=args.measures,
+        repeats=args.repeats,
+    )
+    if args.command == "measure":
+        func = load_function(args.file, args.function)
+        report = measure(
+            func,
+            *sample_args,
+            file=args.file,
+            qualname=args.function,
+            **sample_kwargs,
+            **options,
+        )
+        print(format_table(report.points))
+        print()
+        print(format_metrics(report.comment))
+        return 0
+    report = annotate_file(args.file, args.function, *sample_args, **sample_kwargs, **options)
+    print(format_metrics(report.comment))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())