Move introspect to python/introspect

PiperOrigin-RevId: 728695024
Change-Id: I96433e1e9ee509704d66bd4c1be1906732e4cc55
This commit is contained in:
Alessio Quaglino
2025-02-19 09:00:26 -08:00
committed by Copybara-Service
parent 6a86247810
commit b0e9d08673
19 changed files with 22 additions and 5 deletions
@@ -0,0 +1,149 @@
# Copyright 2022 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utility for formatting AST node as Python code."""
import contextlib
import dataclasses
from typing import Any, Iterable, Mapping, Sequence
INDENT_WIDTH = 4
MAX_LINE_WIDTH = 80
SIMPLE_TYPES = frozenset([int, float, str, bool, bytes, type(None)])
def format_as_python_code(obj: Any) -> str:
"""Formats an AST node object as well-indented Python code."""
formatter = _Formatter()
formatter.add(obj)
return str(formatter)
def _is_all_simple(seq: Iterable[Any]) -> bool:
return all(type(obj) in SIMPLE_TYPES for obj in seq)
class _Formatter:
"""A helper for pretty-printing AST nodes as Python code."""
def __init__(self):
self._line_prefix = ''
self._lines = []
self._add_to_last_line = False
@contextlib.contextmanager
def _indent(self, width: int = INDENT_WIDTH):
self._line_prefix += ' ' * width
yield
self._line_prefix = self._line_prefix[:-width]
@contextlib.contextmanager
def _append_at_end(self, s):
yield
self._lines[-1] += s
def _add_line(self, line: str, no_break: bool = False):
if self._add_to_last_line:
self._lines[-1] += line
else:
self._lines.append(self._line_prefix + line)
self._add_to_last_line = no_break
def _add_dict(self, obj: Mapping[Any, Any]):
"""Adds a dict to the formatted output."""
self._add_line('dict([')
with self._indent():
for k, v in obj.items():
# Try to fit everything into a single line first.
if _is_all_simple((k, v)):
single_line = f'({k!r}, {v!r}),'
if len(self._line_prefix) + len(single_line) <= MAX_LINE_WIDTH:
self._add_line(single_line)
continue
self._add_line(f"('{k}',")
with self._append_at_end('),'):
with self._indent(1):
self.add(v)
self._add_line('])')
def _add_dataclass(self, obj: Any):
"""Adds a dataclass object to the formatted output."""
# Filter out default values.
kv_pairs = []
for k in dataclasses.fields(obj):
v = getattr(obj, k.name)
if v != k.default:
kv_pairs.append((k, v))
# Try to fit everything into a single line first.
if _is_all_simple(v for _, v in kv_pairs):
single_line = ', '.join(f'{k.name}={v!r}' for k, v in kv_pairs)
single_line = f'{obj.__class__.__name__}({single_line})'
if len(self._line_prefix) + len(single_line) <= MAX_LINE_WIDTH:
self._add_line(single_line)
return
self._add_line(obj.__class__.__name__ + '(')
with self._indent():
for k, v in kv_pairs:
self._add_line(k.name + '=', no_break=True)
with self._append_at_end(','):
self.add(v)
self._add_line(')')
def _add_sequence(self, obj: Sequence[Any]) -> None:
"""Adds a sequence to the formatted output."""
default_str = repr(obj)
open_token, close_token = default_str[0], default_str[-1]
# Try to fit everything into a single line first.
if _is_all_simple(obj):
single_line = (
f"{open_token}{', '.join(repr(o) for o in obj)}{close_token}")
if close_token == ')' and len(obj) == 1:
single_line = f'{single_line[:-1]},)'
if len(self._line_prefix) + len(single_line) <= MAX_LINE_WIDTH:
self._add_line(single_line)
return
self._add_line(open_token)
with self._indent():
for v in obj:
with self._append_at_end(','):
self.add(v)
self._add_line(close_token)
def add(self, obj: Any) -> None:
"""Adds an object to the formatted output."""
if _is_all_simple((obj,)):
self._add_line(repr(obj))
elif dataclasses.is_dataclass(obj):
self._add_dataclass(obj)
elif isinstance(obj, Mapping):
self._add_dict(obj)
elif isinstance(obj, Sequence):
self._add_sequence(obj)
else:
raise NotImplementedError
def __str__(self):
lines = []
for line in self._lines:
if len(line) > MAX_LINE_WIDTH:
lines.append(f'{line} # pylint: disable=line-too-long')
else:
lines.append(line)
return '\n'.join(lines)
@@ -0,0 +1,128 @@
# Copyright 2022 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Generates enums.py.
The JSON input can be generated via:
clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h
"""
import json
from typing import Any, Mapping, Sequence
from absl import app
from absl import flags
from introspect import ast_nodes
from . import formatter
_JSON_PATH = flags.DEFINE_string(
'json_path', None,
'Path to the JSON file representing the Clang AST for mujoco.h')
ClangJsonNode = Mapping[str, Any]
def traverse(node, visitor):
visitor.visit(node)
children = node.get('inner', [])
for child in children:
traverse(child, visitor)
class MjEnumVisitor:
"""A Clang AST JSON node visitor for MuJoCo API enum declarations."""
def __init__(self):
self._enums = {}
self._typedefs = {}
def _make_enum(self, node: ClangJsonNode) -> ast_nodes.EnumDecl:
"""Makes a EnumDecl from a Clang AST EnumDecl node."""
name = f"enum {node['name']}"
values = []
for child in node['inner']:
child_kind = child.get('kind')
if child_kind == 'EnumConstantDecl':
next_idx = values[-1][1] + 1 if values else 0
if 'inner' in child:
value = int(child['inner'][0].get('value', next_idx))
else:
value = next_idx
values.append((child['name'], value))
return ast_nodes.EnumDecl(name=name, declname=name, values=dict(values))
def visit(self, node: ClangJsonNode) -> None:
if (node.get('kind') == 'EnumDecl' and
node.get('name', '').startswith('mj')):
enum_decl = self._make_enum(node)
self._enums[enum_decl.name] = enum_decl
elif (node.get('kind') == 'TypedefDecl' and
node['type']['qualType'].startswith('enum mj')):
enum = self._enums[node['type']['qualType']]
self._typedefs[node['name']] = ast_nodes.EnumDecl(
name=node['name'], declname=enum.declname, values=dict(enum.values))
@property
def enums(self) -> Mapping[str, ast_nodes.EnumDecl]:
return self._enums
@property
def typedefs(self) -> Mapping[str, ast_nodes.EnumDecl]:
return self._typedefs
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
with open(_JSON_PATH.value, 'r', encoding='utf-8') as f:
root = json.load(f)
visitor = MjEnumVisitor()
traverse(root, visitor)
enums_str = formatter.format_as_python_code(visitor.typedefs)
print(f'''
# Copyright 2022 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Provides information about MuJoCo API enums.
DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
"""
from typing import Mapping
from .ast_nodes import EnumDecl
ENUMS: Mapping[str, EnumDecl] = {enums_str}
'''.strip()) # `print` adds a trailing newline
if __name__ == '__main__':
app.run(main)
@@ -0,0 +1,159 @@
# Copyright 2022 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Generates functions.py.
The JSON input can be generated via:
clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h
"""
import json
from typing import Any, Mapping, Sequence
from absl import app
from absl import flags
from introspect import ast_nodes
from introspect import type_parsing
from . import formatter
_HEADER_PATH = flags.DEFINE_string(
'header_path', None, 'Path to the original mujoco.h')
_JSON_PATH = flags.DEFINE_string(
'json_path', None,
'Path to the JSON file representing the Clang AST for mujoco.h')
ClangJsonNode = Mapping[str, Any]
def traverse(node, visitor):
visitor.visit(node)
children = node.get('inner', [])
for child in children:
traverse(child, visitor)
class MjFunctionVisitor:
"""A Clang AST JSON node visitor for MuJoCo API function declarations."""
def __init__(self, raw_header):
self._raw_header = raw_header
self._functions = {}
def _make_function(self, node: ClangJsonNode) -> ast_nodes.FunctionDecl:
"""Makes a FunctionDecl from a Clang AST FunctionDecl node."""
name = node['name']
return_type = type_parsing.parse_function_return_type(
node['type']['qualType'])
parameters = []
comments = []
for child in node['inner']:
child_kind = child.get('kind')
if child_kind == 'ParmVarDecl':
parameters.append(self._make_parameter(child))
if child_kind == 'FullComment':
comments.append(self._make_comment(child))
comment = ' '.join(comments).strip()
return ast_nodes.FunctionDecl(
name=name, return_type=return_type, parameters=parameters, doc=comment)
def _make_parameter(
self, node: ClangJsonNode) -> ast_nodes.FunctionParameterDecl:
"""Makes a ParameterDecl from a Clang AST ParmVarDecl node."""
name = node['name']
type_name = node['type']['qualType']
# For a pointer parameters, look up in the original header to see if
# n array extent was declared there.
if type_name.endswith('*'):
decl_begin = node['range']['begin']['offset']
decl_end = node['range']['end']['offset'] + node['range']['end']['tokLen']
decl = self._raw_header[decl_begin:decl_end]
name_begin = node['loc']['offset'] - decl_begin
name_end = name_begin + node['loc']['tokLen']
type_name = decl[:name_begin] + decl[name_end:]
return ast_nodes.FunctionParameterDecl(
name=name, type=type_parsing.parse_type(type_name))
def _make_comment(self, node: ClangJsonNode) -> str:
"""Makes a comment string from a Clang AST FullComment node."""
kind = node.get('kind')
if kind == 'TextComment':
return node['text'].replace('\N{NO-BREAK SPACE}', '&nbsp;')
else:
strings = []
for child in node['inner']:
strings.append(self._make_comment(child))
return ''.join(strings)
def visit(self, node: ClangJsonNode) -> None:
if (node.get('kind') == 'FunctionDecl' and
node.get('name', '').startswith('mj')):
func_decl = self._make_function(node)
self._functions[func_decl.name] = func_decl
@property
def functions(self) -> Mapping[str, ast_nodes.FunctionDecl]:
return self._functions
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
with open(_JSON_PATH.value, 'r', encoding='utf-8') as f:
root = json.load(f)
with open(_HEADER_PATH.value, 'r') as f:
visitor = MjFunctionVisitor(f.read())
traverse(root, visitor)
functions_str = formatter.format_as_python_code(visitor.functions)
print(f'''
# Copyright 2022 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Provides information about MuJoCo API functions.
DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
"""
from typing import Mapping
from .ast_nodes import ArrayType
from .ast_nodes import FunctionDecl
from .ast_nodes import FunctionParameterDecl
from .ast_nodes import PointerType
from .ast_nodes import ValueType
FUNCTIONS: Mapping[str, FunctionDecl] = {functions_str}
'''.strip()) # `print` adds a trailing newline
if __name__ == '__main__':
app.run(main)
@@ -0,0 +1,276 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Generates structs.py.
The JSON input can be generated via:
clang -Xclang -ast-dump=json -fsyntax-only -fparse-all-comments -x c mujoco.h
"""
import itertools
import json
import os
import re
from typing import Any, Mapping, Sequence, Union
from absl import app
from absl import flags
from introspect import ast_nodes
from introspect import type_parsing
from . import formatter
_JSON_PATH = flags.DEFINE_string(
'json_path', None,
'Path to the JSON file representing the Clang AST for mujoco.h')
ClangJsonNode = Mapping[str, Any]
_ANONYMOUS_KEY_PATTERN = re.compile(r'\d+:\d+(?=\))')
_EXCLUDED = (
'mjpPlugin',
'mjpPlugin_',
'mjpResourceProvider',
'mjpResourceProvider_',
'mjResource',
'mjResource_',
)
_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s\s+\((.+) x (.+)\)\Z')
def traverse(node, visitor):
visitor.visit(node)
children = node.get('inner', [])
for child in children:
traverse(child, visitor)
class _AnonymousTypePlaceholder(ast_nodes.ValueType):
def __init__(self, anonymous_key: str):
self.name = anonymous_key
self.is_const = False
self.is_volatile = False
class MjStructVisitor:
"""A Clang AST JSON node visitor for MuJoCo API struct declarations."""
def __init__(self):
self._structs = {}
self._anonymous = {}
self._typedefs = {}
def _normalize_type(
self, declname: str
) -> Union[
ast_nodes.ValueType, ast_nodes.PointerType, ast_nodes.ArrayType]:
"""Resolves anonymous structs/unions and looks up existing typedefs."""
# Check for anonymous struct/union.
if '(unnamed ' in declname:
m = _ANONYMOUS_KEY_PATTERN.search(declname)
if not m:
raise RuntimeError('cannot parse anonymous key from {m!r}')
return _AnonymousTypePlaceholder(m.group(0))
# Lookup typedef name and use it instead if one exists.
for k, v in self._typedefs.items():
if declname == v.declname:
return type_parsing.parse_type(k)
# No valid normalization, just parse the declname.
return type_parsing.parse_type(declname)
def _make_comment(self, node: ClangJsonNode) -> str:
"""Makes a comment string from a Clang AST FullComment node."""
kind = node.get('kind')
if kind == 'TextComment':
return node['text'].replace('\N{NO-BREAK SPACE}', '&nbsp;').strip()
else:
strings = []
for child in node['inner']:
strings.append(self._make_comment(child))
return ''.join(strings).strip()
def _make_field(
self, node: ClangJsonNode
) -> Union[ast_nodes.StructFieldDecl, _AnonymousTypePlaceholder]:
"""Makes a StructFieldDecl object from a Clang AST FieldDecl node."""
doc = ''
for child in node.get('inner', ()):
if child['kind'] == 'FullComment':
doc = self._make_comment(child)
if 'name' in node:
field_type = self._normalize_type(node['type']['qualType'])
m = _ARRAY_COMMENT_PATTERN.match(doc)
if m is None:
array_extent = None
else:
doc = m.group(1)
array_extent_0 = m.group(2)
array_extent_1 = m.group(3)
try:
array_extent_1 = int(array_extent_1)
except ValueError:
pass
if array_extent_1 == 1:
array_extent = (array_extent_0,)
else:
array_extent = (array_extent_0, array_extent_1)
return ast_nodes.StructFieldDecl(
name=node['name'], type=field_type, doc=doc,
array_extent=array_extent)
else:
return _AnonymousTypePlaceholder(self._make_anonymous_key(node))
def _make_struct(
self, node: ClangJsonNode
) -> Union[ast_nodes.AnonymousStructDecl, ast_nodes.StructDecl]:
"""Makes a Decl object from a Clang AST RecordDecl node."""
name = f"{node['tagUsed']} {node['name']}" if 'name' in node else ''
fields = []
for child in node.get('inner', ()):
child_kind = child.get('kind')
if child_kind == 'FieldDecl':
fields.append(self._make_field(child))
if name:
return ast_nodes.StructDecl(name=name, declname=name, fields=fields)
elif node['tagUsed'] == 'union':
return ast_nodes.AnonymousUnionDecl(fields=fields)
else:
return ast_nodes.AnonymousStructDecl(fields=fields)
def _is_mujoco_type(self, node: ClangJsonNode) -> bool:
node_name = node.get('name', '')
included_from = os.path.basename(
node['loc'].get('includedFrom', {}).get('file', '')
)
return node_name not in _EXCLUDED and (
node_name.startswith('mj')
or included_from == 'mujoco.h'
or included_from.startswith('mj')
)
def _make_anonymous_key(self, node: ClangJsonNode) -> str:
line = node['loc']['line']
col = node['loc']['col']
return f'{line}:{col}'
def visit(self, node: ClangJsonNode) -> None:
"""Visits a JSON node."""
if node.get('kind') == 'RecordDecl' and self._is_mujoco_type(node):
struct_decl = self._make_struct(node)
if hasattr(struct_decl, 'name'):
self._structs[struct_decl.name] = struct_decl
else:
anonymous_key = self._make_anonymous_key(node)
if anonymous_key in self._anonymous:
raise RuntimeError(
f'duplicate key for anonymous struct: {anonymous_key}')
self._anonymous[anonymous_key] = struct_decl
elif (node.get('kind') == 'TypedefDecl' and
node['type']['qualType'].startswith('struct mj') and
node['name'] not in _EXCLUDED):
declname = node['type']['qualType']
try:
struct = self._structs[declname]
except KeyError:
self._typedefs[node['name']] = ast_nodes.StructDecl(
name=node['name'], declname=declname, fields=())
else:
self._typedefs[node['name']] = ast_nodes.StructDecl(
name=node['name'], declname=struct.declname, fields=struct.fields)
def resolve_all_anonymous(self) -> None:
"""Replaces anonymous struct placeholders with corresponding decl."""
for struct in itertools.chain(
self._structs.values(), self._typedefs.values()
):
fields = []
for field in struct.fields:
if isinstance(field, _AnonymousTypePlaceholder):
fields.append(self._anonymous[field.name])
elif isinstance(field.type, _AnonymousTypePlaceholder):
fields.append(
ast_nodes.StructFieldDecl(
name=field.name,
type=self._anonymous[field.type.name],
doc=field.doc,
)
)
else:
fields.append(field)
struct.fields = tuple(fields)
@property
def structs(self) -> Mapping[str, ast_nodes.StructDecl]:
return self._structs
@property
def typedefs(self) -> Mapping[str, ast_nodes.StructDecl]:
return self._typedefs
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
with open(_JSON_PATH.value, 'r', encoding='utf-8') as f:
root = json.load(f)
visitor = MjStructVisitor()
traverse(root, visitor)
visitor.resolve_all_anonymous()
structs_str = formatter.format_as_python_code(visitor.typedefs)
print(f'''
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Provides information about MuJoCo API structs.
DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
"""
from typing import Mapping
from .ast_nodes import AnonymousStructDecl
from .ast_nodes import AnonymousUnionDecl
from .ast_nodes import ArrayType
from .ast_nodes import PointerType
from .ast_nodes import StructDecl
from .ast_nodes import StructFieldDecl
from .ast_nodes import ValueType
STRUCTS: Mapping[str, StructDecl] = {structs_str}
'''.strip()) # `print` adds a trailing newline
if __name__ == '__main__':
app.run(main)