1
"""This module provides a wrapper that will work for any OMF data object or
2
project files.
3

4
Example Use
5
-----------
6

7
Use the wrapper provided in ``omfvista`` to wrap any ``omf`` data object:
8

9
.. code-block:: python
10

11
    import omfvista
12

13
    omfvista.wrap(data)
14

15

16
Here's an example using the sample data hosted in the `OMF repository`_.
17

18
.. _OMF repository: https://github.com/gmggroup/omf/tree/master/assets
19

20
.. code-block:: python
21

22
    import omf
23
    import omfvista
24

25
    # Read all elements
26
    reader = omf.OMFReader('test_file.omf')
27
    project = reader.get_project()
28

29
    # Iterate over the elements and add converted VTK objects to dictionary:
30
    data = dict()
31
    for e in project.elements:
32
        d = omfvista.wrap(e)
33
        data[e.name] = d
34

35
Or better yet, just use the project loader:
36

37
.. code-block:: python
38

39
    import omfvista
40
    data = omfvista.load_project('test_file.omf')
41

42
"""
43

44

45 16
__all__ = [
46
    'wrap',
47
    'project_to_vtk',
48
    'load_project',
49
]
50

51 16
__displayname__ = 'Wrapper'
52

53 16
import omf
54 16
import omfvista
55 16
import pyvista
56

57

58 16
def wrap(data):
59
    """Wraps the OMF data object/project as a VTK data object. This is the
60
    primary function that an end user will harness.
61

62
    Args:
63
        data: any OMF data object
64

65
    Example:
66
        >>> import omf
67
        >>> import omfvista
68

69
        >>> # Read all elements
70
        >>> reader = omf.OMFReader('test_file.omf')
71
        >>> project = reader.get_project()
72

73
        >>> # Iterate over the elements and add converted VTK objects to dictionary:
74
        >>> data = dict()
75
        >>> for e in project.elements:
76
        >>>     d = omfvista.wrap(e)
77
        >>>     data[e.name] = d
78

79
    """
80
    # Allow recursion
81 16
    if isinstance(data, (list, tuple)):
82 16
        multi = pyvista.MultiBlock()
83 16
        for i, item in enumerate(data):
84 16
            multi.append(wrap(item))
85 14
            multi.set_block_name(i, item.name)
86 14
        return multi
87
    # Define wrappers
88 16
    wrappers = {
89
        'LineSetElement': omfvista.line_set_to_vtk,
90
        'PointSetElement': omfvista.point_set_to_vtk,
91
        # Surfaces
92
        'SurfaceGeometry': omfvista.surface_geom_to_vtk,
93
        'SurfaceGridGeometry': omfvista.surface_grid_geom_to_vtk,
94
        'SurfaceElement': omfvista.surface_to_vtk,
95
        # Volumes
96
        'VolumeGridGeometry': omfvista.volume_grid_geom_to_vtk,
97
        'VolumeElement': omfvista.volume_to_vtk,
98
        'Project': omfvista.project_to_vtk,
99
    }
100
    # get the class name
101 16
    key = data.__class__.__name__
102 16
    try:
103 16
        return wrappers[key](data)
104 2
    except KeyError:
105 0
        raise RuntimeError('Data of type ({}) is not supported currently.'.format(key))
106

107

108 16
def project_to_vtk(project):
109
    """Converts an OMF project (:class:`omf.base.Project`) to a
110
    :class:`pyvista.MultiBlock` data boject
111
    """
112
    # Iterate over the elements and add converted VTK objects a MultiBlock
113 16
    data = pyvista.MultiBlock()
114 16
    for i, e in enumerate(project.elements):
115 16
        d = omfvista.wrap(e)
116 14
        data[i, e.name] = d
117 14
    return data
118

119

120 16
def load_project(filename):
121
    """Loads an OMF project file into a :class:`pyvista.MultiBlock` dataset"""
122 16
    reader = omf.OMFReader(filename)
123 16
    project = reader.get_project()
124 16
    return project_to_vtk(project)
125

126

127
# Now set up the display names for the docs
128 16
load_project.__displayname__ = 'Load Project File'
129 16
project_to_vtk.__displayname__ = 'Project to VTK'
130 16
wrap.__displayname__ = 'The Wrapper'

Read our documentation on viewing source code .

Loading