Skip to content

Commit 93f1b6d

Browse files
XuehaiPanXuehaiPan
committed
chore: dump build time meta-information to C extension
1 parent a8faff8 commit 93f1b6d

6 files changed

Lines changed: 100 additions & 70 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,10 @@ if(NOT DEFINED OPTREE_CXX_WERROR AND NOT "$ENV{OPTREE_CXX_WERROR}" STREQUAL "")
9797
endif()
9898

9999
if(OPTREE_CXX_WERROR)
100-
message(WARNING "Treats all compiler warnings as errors. Set `OPTREE_CXX_WERROR=OFF` to disable this.")
100+
message(
101+
AUTHOR_WARNING
102+
"Treats all compiler warnings as errors. Set `OPTREE_CXX_WERROR=OFF` to disable this."
103+
)
101104
if(MSVC)
102105
string(APPEND CMAKE_CXX_FLAGS " /WX")
103106
else()

‎Makefile‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,8 @@ addlicense-install: go-install
146146
pytest test: pytest-install
147147
$(PYTEST) --version
148148
cd tests && $(PYTHON) -X dev -Walways -Werror -c 'import $(PROJECT_NAME)' && \
149-
$(PYTHON) -X dev -Walways -Werror -c 'import $(PROJECT_NAME)._C; print(f"GLIBCXX_USE_CXX11_ABI={$(PROJECT_NAME)._C.GLIBCXX_USE_CXX11_ABI}")' && \
149+
$(PYTHON) -X dev -Walways -Werror -c \
150+
'import $(PROJECT_NAME)._C; print(f"GLIBCXX_USE_CXX11_ABI={$(PROJECT_NAME)._C.GLIBCXX_USE_CXX11_ABI}")' && \
150151
$(PYTEST) --verbose --color=yes --durations=10 --showlocals \
151152
--cov="$(PROJECT_NAME)" --cov-report=term-missing \
152153
$(PYTESTOPTS) .

‎include/optree/treespec.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,11 @@ using ssize_t = py::ssize_t;
4242

4343
// The maximum depth of a pytree.
4444
#ifndef Py_C_RECURSION_LIMIT
45+
#ifndef Py_DEBUG
4546
#define Py_C_RECURSION_LIMIT 1000
47+
#else
48+
#define Py_C_RECURSION_LIMIT 500
49+
#endif
4650
#endif
4751
#if !defined(PYPY_VERSION) && !(defined(MS_WINDOWS) && defined(Py_DEBUG))
4852
constexpr ssize_t MAX_RECURSION_DEPTH = std::min(1000, Py_C_RECURSION_LIMIT);

‎optree/_C.pyi‎

Lines changed: 76 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@
1717

1818
import builtins
1919
import enum
20+
import platform
2021
from collections.abc import Callable, Collection, Iterable, Iterator
21-
from typing import Any
22+
from typing import Any, Final, final
2223
from typing_extensions import Self
2324

2425
from optree.typing import (
@@ -32,66 +33,23 @@ from optree.typing import (
3233
UnflattenFunc,
3334
)
3435

35-
class InternalError(RuntimeError): ...
36-
37-
MAX_RECURSION_DEPTH: int
38-
3936
# Set if the type allows subclassing (see CPython's Include/object.h)
40-
Py_TPFLAGS_BASETYPE: int # (1UL << 10)
37+
Py_TPFLAGS_BASETYPE: Final[int] # (1UL << 10)
4138

42-
GLIBCXX_USE_CXX11_ABI: bool
39+
# Meta information during build
40+
PY_VERSION: Final[str]
41+
PY_VERSION_HEX: Final[int]
42+
if platform.python_implementation() == 'PyPy': # noqa: PYI002
43+
PYPY_VERSION: Final[str]
44+
PYPY_VERSION_NUM: Final[int]
45+
PYBIND11_VERSION_HEX: Final[int]
46+
PYBIND11_INTERNALS_VERSION: Final[int]
47+
GLIBCXX_USE_CXX11_ABI: Final[bool]
4348

44-
def flatten(
45-
tree: PyTree[T],
46-
/,
47-
leaf_predicate: Callable[[T], bool] | None = None,
48-
none_is_leaf: bool = False,
49-
namespace: str = '',
50-
) -> tuple[list[T], PyTreeSpec]: ...
51-
def flatten_with_path(
52-
tree: PyTree[T],
53-
/,
54-
leaf_predicate: Callable[[T], bool] | None = None,
55-
none_is_leaf: bool = False,
56-
namespace: str = '',
57-
) -> tuple[list[tuple[Any, ...]], list[T], PyTreeSpec]: ...
58-
def make_leaf(
59-
none_is_leaf: bool = False,
60-
namespace: str = '', # unused
61-
) -> PyTreeSpec: ...
62-
def make_none(
63-
none_is_leaf: bool = False,
64-
namespace: str = '', # unused
65-
) -> PyTreeSpec: ...
66-
def make_from_collection(
67-
collection: Collection[PyTreeSpec],
68-
/,
69-
none_is_leaf: bool = False,
70-
namespace: str = '',
71-
) -> PyTreeSpec: ...
72-
def is_leaf(
73-
obj: T,
74-
/,
75-
leaf_predicate: Callable[[T], bool] | None = None,
76-
none_is_leaf: bool = False,
77-
namespace: str = '',
78-
) -> bool: ...
79-
def all_leaves(
80-
iterable: Iterable[T],
81-
/,
82-
leaf_predicate: Callable[[T], bool] | None = None,
83-
none_is_leaf: bool = False,
84-
namespace: str = '',
85-
) -> bool: ...
86-
def is_namedtuple(obj: object | type, /) -> bool: ...
87-
def is_namedtuple_instance(obj: object, /) -> bool: ...
88-
def is_namedtuple_class(cls: type, /) -> bool: ...
89-
def namedtuple_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ...
90-
def is_structseq(obj: object | type, /) -> bool: ...
91-
def is_structseq_instance(obj: object, /) -> bool: ...
92-
def is_structseq_class(cls: type, /) -> bool: ...
93-
def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ...
49+
@final
50+
class InternalError(RuntimeError): ...
9451

52+
@final
9553
class PyTreeKind(enum.IntEnum):
9654
CUSTOM = 0 # a custom type
9755
LEAF = enum.auto() # an opaque leaf node
@@ -105,6 +63,9 @@ class PyTreeKind(enum.IntEnum):
10563
DEQUE = enum.auto() # a collections.deque
10664
STRUCTSEQUENCE = enum.auto() # a PyStructSequence
10765

66+
MAX_RECURSION_DEPTH: Final[int]
67+
68+
@final
10869
class PyTreeSpec:
10970
num_nodes: int
11071
num_leaves: int
@@ -157,6 +118,7 @@ class PyTreeSpec:
157118
def __hash__(self, /) -> int: ...
158119
def __len__(self, /) -> int: ...
159120

121+
@final
160122
class PyTreeIter(Iterator[T]):
161123
def __init__(
162124
self,
@@ -169,6 +131,63 @@ class PyTreeIter(Iterator[T]):
169131
def __iter__(self, /) -> Self: ...
170132
def __next__(self, /) -> T: ...
171133

134+
# Functions
135+
def flatten(
136+
tree: PyTree[T],
137+
/,
138+
leaf_predicate: Callable[[T], bool] | None = None,
139+
none_is_leaf: bool = False,
140+
namespace: str = '',
141+
) -> tuple[list[T], PyTreeSpec]: ...
142+
def flatten_with_path(
143+
tree: PyTree[T],
144+
/,
145+
leaf_predicate: Callable[[T], bool] | None = None,
146+
none_is_leaf: bool = False,
147+
namespace: str = '',
148+
) -> tuple[list[tuple[Any, ...]], list[T], PyTreeSpec]: ...
149+
150+
# Constructors
151+
def make_leaf(
152+
none_is_leaf: bool = False,
153+
namespace: str = '', # unused
154+
) -> PyTreeSpec: ...
155+
def make_none(
156+
none_is_leaf: bool = False,
157+
namespace: str = '', # unused
158+
) -> PyTreeSpec: ...
159+
def make_from_collection(
160+
collection: Collection[PyTreeSpec],
161+
/,
162+
none_is_leaf: bool = False,
163+
namespace: str = '',
164+
) -> PyTreeSpec: ...
165+
166+
# Utility functions
167+
def is_leaf(
168+
obj: T,
169+
/,
170+
leaf_predicate: Callable[[T], bool] | None = None,
171+
none_is_leaf: bool = False,
172+
namespace: str = '',
173+
) -> bool: ...
174+
def all_leaves(
175+
iterable: Iterable[T],
176+
/,
177+
leaf_predicate: Callable[[T], bool] | None = None,
178+
none_is_leaf: bool = False,
179+
namespace: str = '',
180+
) -> bool: ...
181+
def is_namedtuple(obj: object | type, /) -> bool: ...
182+
def is_namedtuple_instance(obj: object, /) -> bool: ...
183+
def is_namedtuple_class(cls: type, /) -> bool: ...
184+
def namedtuple_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ...
185+
def is_structseq(obj: object | type, /) -> bool: ...
186+
def is_structseq_instance(obj: object, /) -> bool: ...
187+
def is_structseq_class(cls: type, /) -> bool: ...
188+
def structseq_fields(obj: tuple | type[tuple], /) -> tuple[str, ...]: ...
189+
190+
# Registration functions
172191
def register_node(
173192
cls: type[Collection[T]],
174193
/,

‎src/optree.cpp‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,27 @@ void BuildModule(py::module_& mod) { // NOLINT[runtime/references]
4848
GetCxxModule(mod);
4949

5050
mod.doc() = "Optimized PyTree Utilities.";
51-
py::register_local_exception<InternalError>(mod, "InternalError", PyExc_SystemError);
52-
mod.attr("MAX_RECURSION_DEPTH") = py::int_(MAX_RECURSION_DEPTH);
5351
mod.attr("Py_TPFLAGS_BASETYPE") = py::int_(Py_TPFLAGS_BASETYPE);
52+
53+
// Meta information during build
54+
mod.attr("PY_VERSION") = py::str(PY_VERSION);
55+
mod.attr("PY_VERSION_HEX") = py::int_(PY_VERSION_HEX);
56+
#ifdef PYPY_VERSION
57+
mod.attr("PYPY_VERSION") = py::str(PYPY_VERSION);
58+
mod.attr("PYPY_VERSION_NUM") = py::int_(PYPY_VERSION_NUM);
59+
#endif
60+
mod.attr("PYBIND11_VERSION_HEX") = py::int_(PYBIND11_VERSION_HEX);
61+
mod.attr("PYBIND11_INTERNALS_VERSION") = py::int_(PYBIND11_INTERNALS_VERSION);
5462
#ifdef _GLIBCXX_USE_CXX11_ABI
5563
// NOLINTNEXTLINE[modernize-use-bool-literals]
5664
mod.attr("GLIBCXX_USE_CXX11_ABI") = py::bool_(static_cast<bool>(_GLIBCXX_USE_CXX11_ABI));
5765
#else
5866
mod.attr("GLIBCXX_USE_CXX11_ABI") = py::bool_(false);
5967
#endif
6068

69+
py::register_local_exception<InternalError>(mod, "InternalError", PyExc_SystemError);
70+
mod.attr("MAX_RECURSION_DEPTH") = py::int_(MAX_RECURSION_DEPTH);
71+
6172
mod.def("register_node",
6273
&PyTreeTypeRegistry::Register,
6374
"Register a Python type. Extends the set of types that are considered internal nodes "

‎tests/test_ops.py‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -59,15 +59,7 @@ def test_import_no_warnings():
5959
}
6060
assert (
6161
subprocess.check_output(
62-
[
63-
sys.executable,
64-
'-W',
65-
'always',
66-
'-W',
67-
'error',
68-
'-c',
69-
'import optree',
70-
],
62+
[sys.executable, '-Walways', '-Werror', '-c', 'import optree'],
7163
stderr=subprocess.STDOUT,
7264
text=True,
7365
cwd=TEST_ROOT,

0 commit comments

Comments
 (0)