Version 2.1.2: Python bindings, OBJ assets support, bugfixes.
PiperOrigin-RevId: 434731612 Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0
This commit is contained in:
Saran Tunyasuvunakool
@@ -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,126 @@
|
||||
# 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 mujoco.h
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Mapping, Sequence
|
||||
|
||||
from absl import app
|
||||
from absl import flags
|
||||
|
||||
from google3.third_party.mujoco.introspect import ast_nodes
|
||||
from google3.third_party.mujoco.introspect.codegen import formatter
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
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
|
||||
value = int(child['inner'][0].get('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(FLAGS.json_path, '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 google3.third_party.mujoco.introspect.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 mujoco.h
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, Mapping, Sequence
|
||||
|
||||
from absl import app
|
||||
from absl import flags
|
||||
|
||||
from google3.third_party.mujoco.introspect import ast_nodes
|
||||
from google3.third_party.mujoco.introspect import type_parsing
|
||||
from google3.third_party.mujoco.introspect.codegen import formatter
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
flags.DEFINE_string('header_path', None, 'Path to the original mujoco.h')
|
||||
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}', ' ')
|
||||
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(FLAGS.json_path, 'r', encoding='utf-8') as f:
|
||||
root = json.load(f)
|
||||
|
||||
with open(FLAGS.header_path, '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 google3.third_party.mujoco.introspect.ast_nodes import ArrayType
|
||||
from google3.third_party.mujoco.introspect.ast_nodes import FunctionDecl
|
||||
from google3.third_party.mujoco.introspect.ast_nodes import FunctionParameterDecl
|
||||
from google3.third_party.mujoco.introspect.ast_nodes import PointerType
|
||||
from google3.third_party.mujoco.introspect.ast_nodes import ValueType
|
||||
|
||||
FUNCTIONS: Mapping[str, FunctionDecl] = {functions_str}
|
||||
'''.strip()) # `print` adds a trailing newline
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(main)
|
||||
Reference in New Issue
Block a user