tpy-lang 0.3.0.dev0__py3-none-any.whl

tpyc/_data/lib/tpy/json.py

@@ -0,0 +1,221 @@
+# tpy: cpp_namespace("tpystd::json")
+# CPython-compatible json module.
+#
+# Pure-TPy wrapper over tplib.json's JsonReader/JsonWriter. Mirrors the
+# CPython API surface (loads, dumps, JSONDecodeError) for untyped JSON
+# values represented as a recursive union.
+#
+# For typed deserialization into user records, use `tplib.json` with the
+# `@model` decorator -- it's faster and statically typed.
+#
+# Consuming a `loads` result:
+#
+#   loads(s) returns `Own[JsonValue]`. JsonValue is a recursive union, so
+#   `d["k"]` at the use site is rejected -- not every member supports str
+#   subscript, and CPython only gets away with it because its `loads`
+#   returns `Any`. Each level of access needs to be narrowed:
+#
+#       def show(v: JsonValue) -> None:
+#           match v:
+#               case dict() as obj:
+#                   print(obj["users"])
+#               case _:
+#                   print("not a dict")
+#
+#       d = json.loads(s)   # d: Own[JsonValue]
+#       show(d)             # auto-moves into v: JsonValue
+#
+#   For deeply-nested known-shape JSON, prefer `tplib.json` with the
+#   `@model` decorator -- you declare a record class for the shape and
+#   parse directly into it, with no per-level narrowing.
+#
+# Followup (BUGS.md): once `isinstance` accepts type aliases / generic
+# types, exposing `JsonObject = dict[str, JsonValue]` and `JsonList =
+# list[JsonValue]` here will make the natural pattern --
+# `if isinstance(d, JsonObject): d["k"]` -- work without writing a
+# helper or using `match`. Until then the helper-with-match pattern
+# above is the cleanest.
+#
+# Not yet supported (deferred):
+#   - load(fp) / dump(obj, fp): blocked on `io` module
+#   - JSONEncoder / JSONDecoder classes (extension hooks)
+#   - dumps kwargs: ensure_ascii (always behaves as if False -- non-ASCII
+#     is emitted as raw UTF-8), separators (uses CPython defaults),
+#     allow_nan, default, cls, skipkeys
+#   - loads kwargs: object_hook, object_pairs_hook, parse_float, parse_int,
+#     parse_constant
+from tpy import Int32, Own, error_return
+from tplib.json import JsonError, JsonReader, JsonToken, JsonWriter
+
+
+type JsonValue = (
+    None
+    | bool
+    | int
+    | float
+    | str
+    | list[JsonValue]
+    | dict[str, JsonValue]
+)
+
+
+class JSONDecodeError(ValueError):
+    """CPython-compatible JSON decode error.
+
+    Thrown via normal try/except so callers don't need `@error_return`
+    decoration. Carries CPython-named attributes (msg, doc, pos, lineno,
+    colno). Subclasses ValueError to match CPython's hierarchy --
+    `except ValueError` catches it.
+    """
+    msg: str
+    doc: str
+    pos: Int32
+    lineno: Int32
+    colno: Int32
+
+    def __init__(self, msg: str, doc: str, pos: Int32) -> None:
+        # `message` is BaseException's runtime field (used by __str__);
+        # `msg` is CPython's documented attribute on JSONDecodeError.
+        # Always equal -- both names are exposed for compatibility.
+        self.message = msg
+        self.msg = msg
+        self.doc = doc
+        self.pos = pos
+        # 1-based to match CPython's JSONDecodeError.lineno/colno contract.
+        line: Int32 = 1
+        col: Int32 = 1
+        i: Int32 = 0
+        end = pos
+        if end > len(doc):
+            end = len(doc)
+        while i < end:
+            if doc[i] == "\n":
+                line += 1
+                col = 1
+            else:
+                col += 1
+            i += 1
+        self.lineno = line
+        self.colno = col
+
+
+# ----------------------------------------------------------------------
+# loads
+# ----------------------------------------------------------------------
+
+# Internal recursion uses JsonError directly; loads() converts to
+# JSONDecodeError at the boundary. Avoids per-call try/except wrappers
+# around every JsonReader operation. Own[JsonValue] return is required:
+# recursive unions containing reference-type members (list/dict) need
+# explicit ownership transfer to return freshly built values.
+@error_return(JsonError)
+def _read_value(reader: JsonReader) -> Own[JsonValue]:
+    tok = reader.peek()
+    if tok == JsonToken.OBJECT_START:
+        d: dict[str, JsonValue] = {}
+        reader.read_object_start()
+        while reader.has_next():
+            key = reader.read_key()
+            # Intermediate local: dict subscript assignment doesn't accept
+            # an Own[T] rvalue directly today. Revisit if the compiler
+            # learns to auto-move Own into subscript slots.
+            val = _read_value(reader)
+            d[key] = val
+        reader.read_object_end()
+        return d
+    if tok == JsonToken.ARRAY_START:
+        a: list[JsonValue] = []
+        reader.read_array_start()
+        while reader.has_next():
+            item = _read_value(reader)
+            a.append(item)
+        reader.read_array_end()
+        return a
+    if tok == JsonToken.STRING:
+        return reader.read_str()
+    if tok == JsonToken.NUMBER:
+        raw = reader.read_number_raw()
+        # Preserve BigInt precision: parse via int() unless the lexeme has
+        # float syntax. CPython treats `1` and `1.0` as int and float.
+        is_float = False
+        i: Int32 = 0
+        n = len(raw)
+        while i < n:
+            c = raw[i]
+            if c == "." or c == "e" or c == "E":
+                is_float = True
+                break
+            i += 1
+        if is_float:
+            return float(raw)
+        return int(raw)
+    if tok == JsonToken.TRUE:
+        reader.read_bool()
+        return True
+    if tok == JsonToken.FALSE:
+        reader.read_bool()
+        return False
+    if tok == JsonToken.NONE:
+        reader.read_null()
+        return None
+    raise JsonError("Expecting value", reader.position())
+
+
+def loads(s: str) -> Own[JsonValue]:
+    """Deserialize a JSON document string into a JsonValue.
+
+    Returns one of: None, bool, int (BigInt), float, str, list[JsonValue],
+    dict[str, JsonValue]. Raises JSONDecodeError on malformed input or
+    trailing data after the top-level value.
+    """
+    reader = JsonReader(s)
+    try:
+        value = _read_value(reader)
+    except JsonError as e:
+        raise JSONDecodeError(e.message, s, e.pos)
+    # CPython rejects trailing non-whitespace; peek() skips whitespace and
+    # returns END at EOF.
+    if reader.peek() != JsonToken.END:
+        raise JSONDecodeError("Extra data", s, reader.position())
+    return value
+
+
+# ----------------------------------------------------------------------
+# dumps
+# ----------------------------------------------------------------------
+
+def _write_value(v: JsonValue, w: JsonWriter, sort_keys: bool) -> None:
+    match v:
+        case None:
+            w.write_null()
+        case bool() as b:
+            w.write_bool(b)
+        case int() as n:
+            w.write_bigint(n)
+        case float() as f:
+            w.write_float(f)
+        case str() as s:
+            w.write_str(s)
+        case list() as items:
+            w.array_start()
+            for item in items:
+                _write_value(item, w, sort_keys)
+            w.array_end()
+        case dict() as d:
+            w.object_start()
+            keys = sorted(d.keys()) if sort_keys else list(d.keys())
+            for k in keys:
+                w.key(k)
+                _write_value(d[k], w, sort_keys)
+            w.object_end()
+
+
+def dumps(obj: JsonValue, indent: Int32 = 0, sort_keys: bool = False) -> str:
+    """Serialize a JsonValue to a JSON-formatted str.
+
+    indent: non-zero turns on pretty-printing with that many spaces per level.
+    sort_keys: when True, dict keys are emitted in sorted order.
+    """
+    w = JsonWriter(indent)
+    _write_value(obj, w, sort_keys)
+    return w.finish()

tpyc/_data/lib/tpy/math.py

@@ -0,0 +1,406 @@
+# tpy: cpp_namespace("tpystd::math")
+# tpy: include("<tpy/stdlib/math.hpp>")
+#
+# Pending items to reach full CPython `math` parity. Each is blocked on a
+# specific compiler gap tracked in TODO.md / BUGS.md; once fixed, update
+# this module.
+#
+# - Tuple inputs to `prod`, `fsum`, `sumprod`, `dist`. CPython accepts
+#   tuples for all four (`math.prod((1.0, 2.0, 3.0))`). Currently fails
+#   conformance to `Iterable[float]`. Blocked on the "Tuple iteration +
+#   Iterable conformance bundle" TODO.md entry -- needs coordinated sema,
+#   codegen, and runtime work (std::tuple has no begin/end).
+# - Generic over int type for `gcd`, `isqrt` (result bounded by input, so
+#   generic is safe). Currently `int`-only forces BigInt allocation for
+#   small-int callers under --default-int=Int32. Pure-TPy generic is
+#   blocked on two language gaps: the missing `AnyInt` protocol (BigInt +
+#   AnyFixedInt) and the ban on `-x` / `abs(x)` / param reassignment
+#   inside unbounded-generic bodies. See TODO.md "No `AnyInt` protocol
+#   covering BigInt + all `AnyFixedInt`". `lcm`, `factorial`, `perm`,
+#   `comb` should probably stay `int`-only even once that lands -- their
+#   results can exceed fixed-int range (e.g. factorial(13) > Int32::max).
+#
+# Algorithmic follow-ups (correctness under our implementation, but not
+# bit-exact or optimal vs CPython):
+#
+# - `fsum`: currently Neumaier compensated summation. Significantly better
+#   than naive but not bit-exact. CPython uses Shewchuk's partials
+#   algorithm (maintains a growing list of non-overlapping partials,
+#   guarantees the sum rounds correctly to a single double). Worth
+#   adopting if users depend on fsum for bit-exact precision.
+# - `sumprod`: naive fold. CPython uses compensated arithmetic (similar
+#   to fsum) for precision. Large vectors can lose low bits in our impl.
+# - `log(x, base)`: always `log(x) / log(base)`. CPython special-cases
+#   `base == 2` to dispatch to `log2` (one fewer transcendental call).
+#   Minor; consider if benchmarks warrant.
+
+from typing import Final, Iterable, overload
+from tpy import Int32
+from tpy.extern import native, cpp_template, type_param_default, DefaultInt
+
+pi: Final[float] = 3.141592653589793
+tau: Final[float] = 6.283185307179586
+e: Final[float] = 2.718281828459045
+inf: Final[float] = 1e309
+nan: Final[float] = float("nan")
+
+@overload
+@native("std::log")
+def log(x: float) -> float: ...
+
+@overload
+def log(x: float, base: float) -> float:
+    return log(x) / log(base)
+
+@native("std::log10")
+def log10(x: float) -> float: ...
+
+@native("std::log2")
+def log2(x: float) -> float: ...
+
+@native("std::sqrt")
+def sqrt(x: float) -> float: ...
+
+@native("std::cbrt")
+def cbrt(x: float) -> float: ...
+
+@native("std::pow")
+def pow(x: float, y: float) -> float: ...
+
+@native("std::exp")
+def exp(x: float) -> float: ...
+
+@native("std::exp2")
+def exp2(x: float) -> float: ...
+
+@native("std::expm1")
+def expm1(x: float) -> float: ...
+
+@native("std::log1p")
+def log1p(x: float) -> float: ...
+
+@native("std::tgamma")
+def gamma(x: float) -> float: ...
+
+@native("std::lgamma")
+def lgamma(x: float) -> float: ...
+
+@native("std::erf")
+def erf(x: float) -> float: ...
+
+@native("std::erfc")
+def erfc(x: float) -> float: ...
+
+@native("std::nextafter")
+def nextafter(x: float, y: float) -> float: ...
+
+@native("std::ldexp")
+def ldexp(x: float, i: Int32) -> float: ...
+
+@native("std::fma")
+def fma(x: float, y: float, z: float) -> float: ...
+
+@native("tpy::stdlib::math::modf")
+def modf(x: float) -> tuple[float, float]: ...
+
+# Exponent always fits in a 16+ bit signed int (actual range
+# [-1073, 1024] for IEEE 754 double; see tpy/stdlib/math.hpp for why).
+# Defaults to DefaultInt (respects --default-int); user can pick Int32,
+# Int64, BigInt, etc. explicitly. Zero-allocation for fixed-width T;
+# BigInt allocates.
+@type_param_default(T=DefaultInt)
+@cpp_template("::tpy::stdlib::math::frexp<{T}>({0})")
+def frexp[T](x: float) -> tuple[float, T]: ...
+
+@native("tpy::stdlib::math::ulp")
+def ulp(x: float) -> float: ...
+
+@native("tpy::BigInt::from_floor")
+def floor(x: float) -> int: ...
+
+@native("tpy::BigInt::from_ceil")
+def ceil(x: float) -> int: ...
+
+@native("std::sin")
+def sin(x: float) -> float: ...
+
+@native("std::cos")
+def cos(x: float) -> float: ...
+
+@native("std::tan")
+def tan(x: float) -> float: ...
+
+@native("std::fabs")
+def fabs(x: float) -> float: ...
+
+@native("std::hypot")
+def _hypot2(x: float, y: float) -> float: ...
+
+def hypot(*coords: float) -> float:
+    if len(coords) == 0:
+        return 0.0
+    # Hypot-fold is overflow-safe: hypot(hypot(a, b), c) == sqrt(a^2 + b^2 + c^2).
+    result = fabs(coords[0])
+    for c in coords[1:]:
+        result = _hypot2(result, c)
+    return result
+
+@native("std::atan2")
+def atan2(y: float, x: float) -> float: ...
+
+@native("std::asin")
+def asin(x: float) -> float: ...
+
+@native("std::acos")
+def acos(x: float) -> float: ...
+
+@native("std::atan")
+def atan(x: float) -> float: ...
+
+@native("std::sinh")
+def sinh(x: float) -> float: ...
+
+@native("std::cosh")
+def cosh(x: float) -> float: ...
+
+@native("std::tanh")
+def tanh(x: float) -> float: ...
+
+@native("std::asinh")
+def asinh(x: float) -> float: ...
+
+@native("std::acosh")
+def acosh(x: float) -> float: ...
+
+@native("std::atanh")
+def atanh(x: float) -> float: ...
+
+@native("std::isnan")
+def isnan(x: float) -> bool: ...
+
+@native("std::isinf")
+def isinf(x: float) -> bool: ...
+
+@native("std::isfinite")
+def isfinite(x: float) -> bool: ...
+
+@native("std::copysign")
+def copysign(x: float, y: float) -> float: ...
+
+@native("std::fmod")
+def fmod(x: float, y: float) -> float: ...
+
+@native("std::remainder")
+def remainder(x: float, y: float) -> float: ...
+
+def radians(x: float) -> float:
+    return x * (pi / 180.0)
+
+def degrees(x: float) -> float:
+    return x * (180.0 / pi)
+
+def _gcd2(a: int, b: int) -> int:
+    if a < 0:
+        a = -a
+    if b < 0:
+        b = -b
+    while b != 0:
+        t: int = b
+        b = a % b
+        a = t
+    return a
+
+def gcd(*ints: int) -> int:
+    if len(ints) == 0:
+        return 0
+    result = ints[0]
+    if result < 0:
+        result = -result
+    for x in ints[1:]:
+        if result == 1:
+            return 1
+        result = _gcd2(result, x)
+    return result
+
+def lcm(*ints: int) -> int:
+    if len(ints) == 0:
+        return 1
+    result = ints[0]
+    if result < 0:
+        result = -result
+    for x in ints[1:]:
+        if result == 0 or x == 0:
+            return 0
+        # (a // gcd(a, b)) * b keeps the intermediate bounded by max(|a|, |b|);
+        # the naive a*b first would blow up for large BigInts.
+        result = (result // _gcd2(result, x)) * x
+        if result < 0:
+            result = -result
+    return result
+
+def factorial(n: int) -> int:
+    if n < 0:
+        raise ValueError("factorial() not defined for negative values")
+    result: int = 1
+    i: int = 2
+    while i <= n:
+        result = result * i
+        i = i + 1
+    return result
+
+def isqrt(n: int) -> int:
+    if n < 0:
+        raise ValueError("isqrt() argument must be nonnegative")
+    if n == 0:
+        return 0
+    # Newton's method from a bit-length-derived initial guess: for an
+    # n-bit input, sqrt(n) has ~n/2 bits, so 1 << ((bits + 1) // 2) is
+    # already an upper bound within a factor of 2 of the answer (factor
+    # is sqrt(2) for even bit_length, exactly 2 for odd). Convergence
+    # drops from O(log n) iterations (when starting at n) to O(log log n).
+    # The `int(1)` LHS forces BigInt arithmetic so the shift can exceed
+    # Int32 width without overflow-checking.
+    x: int = int(1) << ((n.bit_length() + 1) // 2)
+    y: int = (x + n // x) // 2
+    while y < x:
+        x = y
+        y = (x + n // x) // 2
+    return x
+
+@overload
+def perm(n: int) -> int:
+    return factorial(n)
+
+@overload
+def perm(n: int, k: int) -> int:
+    if n < 0 or k < 0:
+        raise ValueError("perm() arguments must be non-negative")
+    if k > n:
+        return 0
+    result: int = 1
+    i: int = 0
+    while i < k:
+        result = result * (n - i)
+        i = i + 1
+    return result
+
+def comb(n: int, k: int) -> int:
+    if n < 0 or k < 0:
+        raise ValueError("comb() arguments must be non-negative")
+    if k > n:
+        return 0
+    if k > n - k:
+        k = n - k
+    result: int = 1
+    i: int = 0
+    while i < k:
+        result = result * (n - i) // (i + 1)
+        i = i + 1
+    return result
+
+def isclose(a: float, b: float, *, rel_tol: float = 1e-09, abs_tol: float = 0.0) -> bool:
+    if a == b:
+        return True
+    if isinf(a) or isinf(b):
+        return False
+    diff: float = fabs(a - b)
+    max_ab: float = fabs(a)
+    if fabs(b) > max_ab:
+        max_ab = fabs(b)
+    return diff <= abs_tol or diff <= rel_tol * max_ab
+
+@overload
+def prod(iterable: Iterable[Int32], *, start: Int32 = Int32(1)) -> Int32:
+    result: Int32 = start
+    for x in iterable:
+        result = result * x
+    return result
+
+@overload
+def prod(iterable: Iterable[int], *, start: int = 1) -> int:
+    result: int = start
+    for x in iterable:
+        result = result * x
+    return result
+
+@overload
+def prod(iterable: Iterable[float], *, start: float = 1.0) -> float:
+    result: float = start
+    for x in iterable:
+        result = result * x
+    return result
+
+def fsum(iterable: Iterable[float]) -> float:
+    # Neumaier summation: more accurate than naive + compensated for large swings.
+    s: float = 0.0
+    c: float = 0.0
+    for x in iterable:
+        t: float = s + x
+        if fabs(s) >= fabs(x):
+            c = c + ((s - t) + x)
+        else:
+            c = c + ((x - t) + s)
+        s = t
+    return s + c
+
+def sumprod(p: Iterable[float], q: Iterable[float]) -> float:
+    # Manual two-iterator drive with length-mismatch detection; CPython uses
+    # zip(p, q, strict=True) internally. We don't have strict=True yet, so
+    # open-code it.
+    s: float = 0.0
+    ip = iter(p)
+    iq = iter(q)
+    while True:
+        try:
+            a = next(ip)
+        except StopIteration:
+            try:
+                next(iq)
+            except StopIteration:
+                return s
+            raise ValueError("sumprod(): input lengths differ")
+        try:
+            b = next(iq)
+        except StopIteration:
+            raise ValueError("sumprod(): input lengths differ")
+        s = s + a * b
+
+def dist(p: Iterable[float], q: Iterable[float]) -> float:
+    # Fold via hypot to avoid overflow when coordinates are large: the naive
+    # sqrt(sum((pi - qi)**2)) overflows when any (pi - qi)**2 exceeds DBL_MAX.
+    # std::hypot(a, b) is IEEE overflow-safe, and hypot-folding is
+    # mathematically equivalent: hypot(hypot(d0, d1), d2) == sqrt(d0^2 + d1^2 + d2^2).
+    ip = iter(p)
+    iq = iter(q)
+    # Seed with first pair (or return 0.0 for empty inputs).
+    try:
+        p0 = next(ip)
+    except StopIteration:
+        try:
+            next(iq)
+        except StopIteration:
+            return 0.0
+        raise ValueError("dist(): input lengths differ")
+    try:
+        q0 = next(iq)
+    except StopIteration:
+        raise ValueError("dist(): input lengths differ")
+    result: float = fabs(p0 - q0)
+    while True:
+        try:
+            pi = next(ip)
+        except StopIteration:
+            try:
+                next(iq)
+            except StopIteration:
+                return result
+            raise ValueError("dist(): input lengths differ")
+        try:
+            qi = next(iq)
+        except StopIteration:
+            raise ValueError("dist(): input lengths differ")
+        result = _hypot2(result, pi - qi)
+
+@type_param_default(T=DefaultInt)
+@cpp_template("::tpy::from_float_check<{T}>({0})")
+def trunc[T](x: float) -> T: ...