feat(clew): extract Ren'Py ASTs
This commit was merged in pull request #20.
This commit is contained in:
298
src/clew/extract_renpy_ast.py
Normal file
298
src/clew/extract_renpy_ast.py
Normal file
@@ -0,0 +1,298 @@
|
||||
from __future__ import print_function, unicode_literals
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
import types
|
||||
import zlib
|
||||
|
||||
try:
|
||||
basestring
|
||||
except NameError:
|
||||
basestring = str
|
||||
|
||||
try:
|
||||
unicode
|
||||
except NameError:
|
||||
unicode = str
|
||||
|
||||
try:
|
||||
long
|
||||
except NameError:
|
||||
long = int
|
||||
|
||||
|
||||
RPYC2_HEADER = b"RENPY RPC2"
|
||||
SCHEMA = "clew-renpy-ast-v1"
|
||||
|
||||
|
||||
def parse_args(argv):
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="extract-renpy-ast",
|
||||
description="Extract compiled Ren'Py script ASTs as generic JSON.",
|
||||
)
|
||||
parser.add_argument("--game-root", required=True)
|
||||
parser.add_argument("--input", required=True)
|
||||
parser.add_argument("--output", required=True)
|
||||
parser.add_argument("--max-depth", type=int)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def read_rpyc_data(path, slot):
|
||||
with open(path, "rb") as handle:
|
||||
header_data = handle.read(1024)
|
||||
|
||||
if header_data[:len(RPYC2_HEADER)] != RPYC2_HEADER:
|
||||
if slot != 1:
|
||||
return None
|
||||
handle.seek(0)
|
||||
return zlib.decompress(handle.read())
|
||||
|
||||
position = len(RPYC2_HEADER)
|
||||
while True:
|
||||
header_slot, start, length = struct.unpack("III", header_data[position:position + 12])
|
||||
if slot == header_slot:
|
||||
break
|
||||
if header_slot == 0:
|
||||
return None
|
||||
position += 12
|
||||
|
||||
handle.seek(start)
|
||||
return zlib.decompress(handle.read(length))
|
||||
|
||||
|
||||
def load_compiled_script(path):
|
||||
import renpy.script
|
||||
|
||||
for slot in (2, 1):
|
||||
data = read_rpyc_data(path, slot)
|
||||
if data:
|
||||
metadata, roots = renpy.script.loads(data)
|
||||
return slot, metadata, roots
|
||||
|
||||
raise RuntimeError("no loadable slot found in {0}".format(path))
|
||||
|
||||
|
||||
def class_name(value):
|
||||
return value.__class__.__module__ + "." + value.__class__.__name__
|
||||
|
||||
|
||||
def collect_slots(value):
|
||||
slots = []
|
||||
try:
|
||||
classes = value.__class__.mro()
|
||||
except AttributeError:
|
||||
classes = [value.__class__]
|
||||
|
||||
for cls in reversed(classes):
|
||||
cls_slots = getattr(cls, "__slots__", [])
|
||||
if isinstance(cls_slots, basestring):
|
||||
cls_slots = [cls_slots]
|
||||
for slot in cls_slots:
|
||||
if slot not in slots:
|
||||
slots.append(slot)
|
||||
return slots
|
||||
|
||||
|
||||
def safe_text(value):
|
||||
if isinstance(value, unicode):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return value.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return {"type": "bytes", "repr": repr(value)}
|
||||
return value
|
||||
|
||||
|
||||
def is_primitive(value):
|
||||
return value is None or isinstance(value, (bool, int, long, float, basestring))
|
||||
|
||||
|
||||
def serialize_function(value):
|
||||
code = getattr(value, "func_code", None) or getattr(value, "__code__", None)
|
||||
result = {
|
||||
"type": "function",
|
||||
"name": safe_text(getattr(value, "func_name", None) or getattr(value, "__name__", None)),
|
||||
"module": safe_text(getattr(value, "__module__", None)),
|
||||
}
|
||||
if code is not None:
|
||||
result["filename"] = safe_text(getattr(code, "co_filename", None))
|
||||
result["firstlineno"] = getattr(code, "co_firstlineno", None)
|
||||
return result
|
||||
|
||||
|
||||
def serialize_method(value):
|
||||
function = getattr(value, "im_func", None) or getattr(value, "__func__", None)
|
||||
receiver = getattr(value, "im_self", None) or getattr(value, "__self__", None)
|
||||
result = serialize_function(function) if function is not None else {
|
||||
"type": "function",
|
||||
"name": safe_text(getattr(value, "__name__", None)),
|
||||
"module": safe_text(getattr(value, "__module__", None)),
|
||||
}
|
||||
result["type"] = "method"
|
||||
if receiver is not None:
|
||||
result["receiver_class"] = class_name(receiver)
|
||||
return result
|
||||
|
||||
|
||||
def serialize_type(value):
|
||||
return {
|
||||
"type": "python-type",
|
||||
"name": safe_text(getattr(value, "__name__", None)),
|
||||
"module": safe_text(getattr(value, "__module__", None)),
|
||||
}
|
||||
|
||||
|
||||
class Serializer(object):
|
||||
def __init__(self, max_depth):
|
||||
self.max_depth = max_depth
|
||||
self.next_id = 1
|
||||
self.seen = {}
|
||||
self.active_containers = set()
|
||||
|
||||
def serialize(self, value, depth=0):
|
||||
if is_primitive(value):
|
||||
return safe_text(value)
|
||||
|
||||
if self.max_depth is not None and depth >= self.max_depth:
|
||||
return {
|
||||
"type": "truncated",
|
||||
"class": class_name(value),
|
||||
"repr": repr(value),
|
||||
}
|
||||
|
||||
if isinstance(value, (list, tuple)):
|
||||
object_id = id(value)
|
||||
if object_id in self.active_containers:
|
||||
return {"type": "recursion", "class": class_name(value)}
|
||||
self.active_containers.add(object_id)
|
||||
items = [self.serialize(item, depth + 1) for item in value]
|
||||
self.active_containers.remove(object_id)
|
||||
return {"type": "tuple" if isinstance(value, tuple) else "list", "items": items}
|
||||
|
||||
if isinstance(value, dict):
|
||||
object_id = id(value)
|
||||
if object_id in self.active_containers:
|
||||
return {"type": "recursion", "class": class_name(value)}
|
||||
self.active_containers.add(object_id)
|
||||
items = []
|
||||
for key, item in sorted(value.items(), key=lambda pair: repr(pair[0])):
|
||||
items.append({"key": self.serialize(key, depth + 1), "value": self.serialize(item, depth + 1)})
|
||||
self.active_containers.remove(object_id)
|
||||
return {"type": "dict", "items": items}
|
||||
|
||||
if isinstance(value, types.FunctionType) or isinstance(value, types.BuiltinFunctionType):
|
||||
return serialize_function(value)
|
||||
|
||||
if isinstance(value, types.MethodType):
|
||||
return serialize_method(value)
|
||||
|
||||
if isinstance(value, type):
|
||||
return serialize_type(value)
|
||||
|
||||
identity = id(value)
|
||||
if identity in self.seen:
|
||||
return {"type": "ref", "id": self.seen[identity]}
|
||||
|
||||
value_id = self.next_id
|
||||
self.next_id += 1
|
||||
self.seen[identity] = value_id
|
||||
|
||||
attributes = {}
|
||||
for slot in collect_slots(value):
|
||||
if slot.startswith("__"):
|
||||
continue
|
||||
try:
|
||||
slot_value = getattr(value, slot)
|
||||
except AttributeError:
|
||||
continue
|
||||
except Exception as error:
|
||||
attributes[slot] = {
|
||||
"type": "attribute-error",
|
||||
"class": class_name(error),
|
||||
"repr": repr(error),
|
||||
}
|
||||
continue
|
||||
attributes[slot] = self.serialize(slot_value, depth + 1)
|
||||
|
||||
value_dict = getattr(value, "__dict__", None)
|
||||
if value_dict:
|
||||
for key, item in sorted(value_dict.items(), key=lambda pair: repr(pair[0])):
|
||||
if isinstance(key, basestring):
|
||||
key_text = safe_text(key)
|
||||
else:
|
||||
key_text = repr(key)
|
||||
attributes[key_text] = self.serialize(item, depth + 1)
|
||||
|
||||
if not attributes:
|
||||
del self.seen[identity]
|
||||
return {
|
||||
"type": "repr",
|
||||
"class": class_name(value),
|
||||
"repr": repr(value),
|
||||
}
|
||||
|
||||
return {
|
||||
"type": "object",
|
||||
"id": value_id,
|
||||
"class": class_name(value),
|
||||
"attributes": attributes,
|
||||
}
|
||||
|
||||
|
||||
def write_json(path, document):
|
||||
parent = os.path.dirname(path)
|
||||
if parent and not os.path.isdir(parent):
|
||||
os.makedirs(parent)
|
||||
|
||||
data = json.dumps(document, separators=(",", ":"), sort_keys=True)
|
||||
if not isinstance(data, bytes):
|
||||
data = data.encode("utf-8")
|
||||
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(data)
|
||||
handle.write(b"\n")
|
||||
|
||||
|
||||
def import_renpy(game_root):
|
||||
sys.path.insert(0, game_root)
|
||||
|
||||
import renpy
|
||||
renpy.import_all()
|
||||
|
||||
class Dummy(object):
|
||||
pass
|
||||
|
||||
renpy.game.script = Dummy()
|
||||
renpy.game.script.all_pyexpr = None
|
||||
renpy.game.script.record_pycode = False
|
||||
renpy.game.log = Dummy()
|
||||
renpy.game.log.mutated = {}
|
||||
|
||||
return renpy
|
||||
|
||||
|
||||
def main(argv):
|
||||
args = parse_args(argv)
|
||||
renpy = import_renpy(args.game_root)
|
||||
slot, metadata, roots = load_compiled_script(args.input)
|
||||
|
||||
serializer = Serializer(args.max_depth)
|
||||
document = {
|
||||
"schema": SCHEMA,
|
||||
"source": os.path.abspath(args.input),
|
||||
"game_root": os.path.abspath(args.game_root),
|
||||
"slot": slot,
|
||||
"python_version": sys.version,
|
||||
"renpy_version": getattr(renpy, "version", None),
|
||||
"metadata": serializer.serialize(metadata),
|
||||
"roots": serializer.serialize(roots),
|
||||
}
|
||||
write_json(args.output, document)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(sys.argv[1:])
|
||||
Reference in New Issue
Block a user