path: root/io_mesh_stl/stl_utils.py

'''
Import and export STL files

Used as a blender script, it load all the stl files in the scene:

blender -P stl_utils.py -- file1.stl file2.stl file3.stl ...
'''

import struct
import mmap
import contextlib
import itertools

# TODO: endien


@contextlib.contextmanager
def mmap_file(filename):
    '''
    Context manager over the data of an mmap'ed file (Read ONLY).


    Example:

    with mmap_file(filename) as m:
        m.read()
        print m[10:50]
    '''
    with open(filename, 'rb') as file:
        # check http://bugs.python.org/issue8046 to have mmap context
        # manager fixed in python
        map = mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ)
        yield map
        map.close()


class ListDict(dict):
    '''
    Set struct with order.

    You can:
       - insert data into without doubles
       - get the list of data in insertion order with self.list

    Like collections.OrderedDict, but quicker, can be replaced if
    ODict is optimised.
    '''

    def __init__(self):
        dict.__init__(self)
        self.list = []
        self._len = 0

    def add(self, item):
        '''
        Add a value to the Set, return its position in it.
        '''
        value = self.setdefault(item, self._len)
        if value == self._len:
            self.list.append(item)
            self._len += 1

        return value

BINARY_HEADER = 80
BINARY_STRIDE = 12 * 4 + 2

def _is_ascii_file(data):
    '''
    This function returns True if the data represents an ASCII file.

    Please note that a False value does not necessary means that the data
    represents a binary file. It can be a (very *RARE* in real life, but
    can easily be forged) ascii file.
    '''
    size = struct.unpack_from('<I', data, BINARY_HEADER)[0]

    return not data.size() == BINARY_HEADER + 4 + BINARY_STRIDE * size

def _binary_read(data):
    # an stl binary file is
    # - 80 bytes of description
    # - 4 bytes of size (unsigned int)
    # - size triangles :
    #
    #   - 12 bytes of normal
    #   - 9 * 4 bytes of coordinate (3*3 floats)
    #   - 2 bytes of garbage (usually 0)

    # OFFSET for the first byte of coordinate (headers + first normal bytes)
    # STRIDE between each triangle (first normal + coordinates + garbage)
    OFFSET = BINARY_HEADER + 4 + 12

    # read header size, ignore description
    size = struct.unpack_from('<I', data, BINARY_HEADER)[0]
    unpack = struct.Struct('<9f').unpack_from

    for i in range(size):
        # read the points coordinates of each triangle
        pt = unpack(data, OFFSET + BINARY_STRIDE * i)
        yield pt[:3], pt[3:6], pt[6:]


def _ascii_read(data):
    # an stl ascii file is like
    # HEADER: solid some name
    # for each face:
    #
    #     facet normal x y z
    #     outerloop
    #          vertex x y z
    #          vertex x y z
    #          vertex x y z
    #     endloop
    #     endfacet

    # strip header
    data.readline()

    while True:
        # strip facet normal // or end
        data.readline()

        # strip outer loup, in case of ending, break the loop
        if not data.readline():
            break

        yield [tuple(map(float, data.readline().split()[1:]))
               for _ in range(3)]

        # strip facet normalend and outerloop end
        data.readline()
        data.readline()


def _binary_write(filename, faces):
    with open(filename, 'wb') as data:
        # header
        # we write padding at header begginning to avoid to
        # call len(list(faces)) which may be expensive
        data.write(struct.calcsize('<80sI') * b'\0')

        # 3 vertex == 9f
        pack = struct.Struct('<9f').pack
        # pad is to remove normal, we do use them
        pad = b'\0' * struct.calcsize('<3f')

        nb = 0
        for verts in faces:
            # write pad as normal + vertexes + pad as attributes
            data.write(pad + pack(*itertools.chain.from_iterable(verts)))
            data.write(b'\0\0')
            nb += 1

        # header, with correct value now
        data.seek(0)
        data.write(struct.pack('<80sI', b"Exported from blender", nb))


def _ascii_write(filename, faces):
    with open(filename, 'w') as data:
        data.write('solid Exported from blender\n')

        for face in faces:
            data.write('''facet normal 0 0 0\nouter loop\n''')
            for vert in face:
                data.write('vertex %f %f %f\n' % vert)
            data.write('endloop\nendfacet\n')

        data.write('endsolid Exported from blender\n')


def write_stl(filename, faces, ascii=False):
    '''
    Write a stl file from faces,

    filename
       output filename

    faces
       iterable of tuple of 3 vertex, vertex is tuple of 3 coordinates as float

    ascii
       save the file in ascii format (very huge)
    '''
    (_ascii_write if ascii else _binary_write)(filename, faces)


def read_stl(filename):
    '''
    Return the triangles and points of an stl binary file.

    Please note that this process can take lot of time if the file is
    huge (~1m30 for a 1 Go stl file on an quad core i7).

    - returns a tuple(triangles, points).

      triangles
          A list of triangles, each triangle as a tuple of 3 index of
          point in *points*.

      points
          An indexed list of points, each point is a tuple of 3 float
          (xyz).

    Example of use:

       >>> tris, pts = read_stl(filename, lambda x:)
       >>> pts = list(pts)
       >>>
       >>> # print the coordinate of the triangle n
       >>> print(pts[i] for i in tris[n])
    '''

    tris, pts = [], ListDict()

    with mmap_file(filename) as data:
        # check for ascii or binary
        gen = _ascii_read if _is_ascii_file(data) else _binary_read

        for pt in gen(data):
            # Add the triangle and the point.
            # If the point is allready in the list of points, the
            # index returned by pts.add() will be the one from the
            # first equal point inserted.
            tris.append([pts.add(p) for p in pt])

    return tris, pts.list


if __name__ == '__main__':
    import sys
    import bpy
    from io_mesh_stl import blender_utils

    filenames = sys.argv[sys.argv.index('--') + 1:]

    for filename in filenames:
        objName = bpy.path.display_name(filename)
        tris, pts = read_stl(filename)

        blender_utils.create_and_link_mesh(objName, tris, pts)