model_generator2/obj.py

from dataclasses import dataclass
import sys
import os

@dataclass(frozen=True)
class FixedPoint:
    s: int # sign
    w: int # whole
    f: int # fraction
    b: int # base

    def to_string(self):
        s_s = "-" if self.s < 0 else ""
        w_s = f"{self.w}"
        f_s = f"{self.f}".rjust(self.b, '0')
        return f"{s_s}{w_s}.{f_s}"

@dataclass(frozen=True)
class VertexPosition:
    x: FixedPoint
    y: FixedPoint
    z: FixedPoint

@dataclass(frozen=True)
class VertexNormal:
    x: FixedPoint
    y: FixedPoint
    z: FixedPoint

@dataclass(frozen=True)
class VertexTexture:
    x: FixedPoint
    y: FixedPoint
    z: FixedPoint

@dataclass(frozen=True)
class ObjectEvent:
    name: str

@dataclass(frozen=True)
class FacePTN:
    position: int
    texture: int
    normal: int

@dataclass
class Face:
    ptn: list[FacePTN]

@dataclass
class Material:
    name: str

@dataclass
class Object:
    name: str
    faces: list[Face]
    material: Material

    def __init__(self, name, faces=None, material=None):
        self.name = name
        self.faces = [] if faces is None else faces
        self.material = material

@dataclass
class ObjFile:
    position: list[VertexPosition]
    normal: list[VertexNormal]
    texture: list[VertexTexture]
    objects: list[Object]

    def __init__(self, position=None, normal=None, texture=None, objects=None):
        self.position = [] if position is None else position
        self.normal = [] if normal is None else normal
        self.texture = [] if texture is None else texture
        self.objects = [] if objects is None else objects

def parse_float(s):
    sign = -1 if s.startswith("-") else 1
    s = s.removeprefix("-")
    if '.' not in s:
        return sign * int(s, 10)
    i, f = s.split('.')
    f_digits = len(f)
    i = int(i, 10)
    f = int(f, 10)

    return FixedPoint(sign, i, f, f_digits)

def parse_vertex(line, n, type):
    assert len(line) == n
    vs = [parse_float(line[i]) for i in range(n)]
    if len(vs) < 3:
        vs.append(0)
    return type(*vs)

def parse_vertex_position(line):
    return parse_vertex(line, 3, VertexPosition)

def parse_vertex_normal(line):
    return parse_vertex(line, 3, VertexNormal)

def parse_vertex_texture(line):
    return parse_vertex(line, 2, VertexTexture)

def parse_face_indices(indices):
    assert "/" in indices
    indices = indices.split("/")
    assert len(indices) == 3, indices
    def face_ix(s):
        i = int(s, 10)
        assert i >= 1
        return i - 1
    return FacePTN(*(face_ix(i) for i in indices))

def parse_face(line):
    return Face([parse_face_indices(indices) for indices in line])

def parse_object_event(line):
    assert len(line) == 1
    name, = line
    name = name.replace(".", "_").replace("-", "_")
    return ObjectEvent(name)

def parse_usemtl(line):
    assert len(line) == 1
    name, = line
    return Material(name)

def parse_line(line):
    t, *line = line.split(' ')
    if t == '#':
        return None
    if t == 'usemtl':
        return parse_usemtl(line)
    if t == 'mtllib':
        return None
    if t == 'o':
        return parse_object_event(line)
    if t == 'v':
        return parse_vertex_position(line)
    if t == 'vn':
        return parse_vertex_normal(line)
    if t == 'vt':
        return parse_vertex_texture(line)
    if t == 'f':
        return parse_face(line)
    if t == 's':
        # smooth shading
        return None
    if t == 'l':
        # polyline
        return None
    assert False, (t, line)

def parse_obj_lines(lines):
    file = ObjFile()
    object = Object(None)
    for line in lines:
        x = parse_line(line)
        if x is None:
            continue
        elif type(x) is VertexPosition:
            file.position.append(x)
        elif type(x) is VertexNormal:
            file.normal.append(x)
        elif type(x) is VertexTexture:
            file.texture.append(x)
        elif type(x) is ObjectEvent:
            if object.faces:
                assert object.name != None
                file.objects.append(object)
            object = Object(x.name.replace('-', '_').replace('.', '_'))
        elif type(x) is Face:
            object.faces.append(x)
        elif type(x) is Material:
            if object.material != None:
                file.objects.append(object)
                object = Object(object.name + "_mtl_" + x.name)
            object.material = x
        else:
            assert False, x
    if object.faces:
        assert object.name != None
        file.objects.append(object)
    return file


def parse_obj_file(filename):
    with open(filename, "r") as f:
        lines = f.read().strip().split("\n")
    file = parse_obj_lines(lines)
    return file

if __name__ == "__main__":
    file = parse_obj_file(sys.argv[1])
    from pprint import pprint
    pprint(file)