ma_parse 2日目

ma_parse.py

# -*- coding: utf-8 -*-
import re

# constant
TRUE_KEYWORDS = ["true", "on", "yes"]
FALSE_KEYWORDS = ["false", "off", "no"]
MA_COMMAND = ["requires",
              "fileInfo",
              "currentUnit",
              "createNode",
              "setAttr",
              "addAttr",
              "connectAttr",
              "disconnectAttr",
              "parent",
              "select",
              "lockNode",
              "file",
              "relationship"]


class Scene(object):

    __nnode = {}
    __reference = {}
    set_node = None

    def __init__(self, path):
        self.path = path
        # init
        self.fileInfo = FileInfo()
        self.__node = {}
        # ---- #

    @property
    def node(self):
        return self.__node

    @property
    def reference(self):
        return self.__reference

    def get_node(self, t=None):
        if t is None:
            return self.__node
        else:
            return [self.__node[x] for x in self.__node.keys() if t == self.__node[x].node_type]

    def parse(self):
        """
        １行ずつ解析して、結果を変数に入れる
        """
        for i in self._read():
            cmd_type = self._get_command_type(i)
            if cmd_type == "fileInfo":
                self.on_fileinfo(i)
                self.set_node = None
            elif cmd_type == "createNode":
                self.on_createnode(i)
            elif cmd_type == "setAttr":
                self.on_setattr(i)
            elif cmd_type == "file":
                self.on_file(i)
            else:
                self.set_node = None
            self.before_cmd = cmd_type

    def on_setattr(self, line):

        a = Attribute()
        a.parse(line)

        if self.set_node is not None:
            self.__node[self.set_node].add_attribute(a.attr_name, a)

    def on_file(self, line):

        line = re.sub("file \"", "", line)
        buff = file_parse(line)
        self.reference[buff[0]] = Scene(line[0])

    def on_fileinfo(self, line):
        """
        fileInfoコマンドの解析
        """
        line = re.sub("fileInfo \"", "", line)
        line = re.sub(";", "", line)
        line = re.sub("\"$", "", line)
        buff = line.split("\" \"")
        self.fileInfo.set_value(buff[0], buff[1])

    def on_createnode(self, line):
        line = re.sub("createNode ", "", line)
        line = re.sub(";", "", line)
        buff = line.split(" ")
        if buff[0] != "setAttr":
            node_type, option = createnode_parse(buff)
            node = Node(node_type, option)
            self.node[node.name] = node
            self.set_node = node.name

    def _read(self):

        with open(self.path, 'r') as f:
            lines = f.readlines()
        buff = []
        for line in lines:
            line = line.strip()
            if line[0:2] == "//":
                pass
            else:
                if line[-1] == ";":
                    buff.append(line)
                    yield "".join(buff)
                    buff = []
                else:
                    buff.append(line)

    def _get_command_type(self, line):

        for cmd in MA_COMMAND:
            if re.search("{0}".format(cmd), line) is not None:
                return cmd


class FileInfo:

    application = ""
    product = ""
    version = ""
    cutIdentifer = ""
    os = ""

    def set_value(self, attr, value):
        if attr == "application":
            self.application = value
        elif attr == "product":
            self.product = value
        elif attr == "version":
            self.version = value
        elif attr == "cutIdentifer":
            self.cutIdentifer = value
        elif attr == "osv":
            self.os = value


class Node:

    name = ""
    node_type = ""
    attribute = {}

    def __init__(self, node_type, opt):
        # init
        self.attribute = {}
        # ---- #
        self.node_type = node_type
        if "-n" in opt:
            self.name = opt['-n']

    def __str__(self):
        return self.name

    def __unicode(self):
        return self.name

    def add_attribute(self, attr, value):
        self.attribute[attr] = value


class Attribute:

    attr_name = None

    def __init__(self):
        pass

    def __str__(self):
        return self.attr_name

    def __unicode__(self):
        return self.attr_name

    def parse(self, line):

        buff = attribute_parse(line)

        self.attr_name = strip_dot(buff[0])
        self.opt = buff[1]
        self.vals = buff[2]

    def get(self):
        if "-type" in self.opt:
            n_type = self.opt['-type']
            if n_type == "double3":
                return [float(x) for x in self.vals]
            elif n_type == "string":
                return self.vals[0]
            else:
                if len(self.vals[0]) == 1:
                    return self.vals[0]
        else:
            if len(self.vals) == 1:
                if self.vals[0] in TRUE_KEYWORDS:
                    return True
                elif self.vals[0] in FALSE_KEYWORDS:
                    return False
                return self.vals[0]


# ------------------- #
# sub_functions
# ------------------- #

# ------------- #
# コマンドのパース
# ------------- #
def file_parse(line):
    """
    fileコマンドをパースする
    """
    line = re.sub(";", "", line)
    buff = line.split(" ")
    buff.pop(0)

    flg = None
    opt = {}
    vals = strip_quotes(buff.pop(-1))

    for i in buff:
        if i != "":
            if flg is None:
                flg = i
            else:
                opt[flg] = strip_quotes(i)
                flg = None
    opt = opt
    vals = vals
    return vals, opt


def attribute_parse(line):

    line = re.sub(";", "", line)
    buff = line.split(" ")
    buff.pop(0)

    flg = None
    opt = {}
    vals = []
    attr_name = None

    for i in buff:
        if i != "":
            if is_flag(i) is False and flg is None:
                if attr_name is None:
                    attr_name = strip_quotes(i)
                else:
                    vals.append(strip_quotes(i))
            else:
                if flg is None:
                    flg = i
                else:
                    opt[flg] = strip_quotes(i)
                    flg = None
    return attr_name, opt, vals

def strip_quotes(s):
    if s[0] == '"':
        return s[1:-1]
    else:
        return s


def strip_dot(s):
    if s[0] == ".":
        return s[1:]
    else:
        return s


def createnode_parse(buff):
    node_type = buff.pop(0)
    opt = {}
    flg = None
    for i in buff:
        if flg is not None:
            if is_flag(i) is True:
                opt[flg] = True
                flg = i
            else:
                opt[flg] = strip_quotes(i)
                flg = None
        else:
            flg = i
    return node_type, opt


def is_flag(s):
    if is_num(s) is True:
        return False
    else:
        if s[0] == "-":
            return True
        else:
            return False


def is_num(s):
    try:
        float(s)
        return True
    except:
        return False


# -------------------------- #
# test
# -------------------------- #
a = Scene("i:/dev/sample_scn/ma_parse/scenes/ref.ma")
a.parse()
# referenceを取得
print a.reference
# 引数で指定したノードをGetして、アトリビュート一覧を取得する
print a.get_node("camera")[0].attribute["den"].get()