1import os
2
3from .. import util
4from .. import visual
5
6from ..base import Trimesh
7from ..parent import Geometry
8from ..points import PointCloud
9from ..scene.scene import Scene, append_scenes
10from ..constants import log_time, log
11
12from . import misc
13from .ply import _ply_loaders
14from .stl import _stl_loaders
15from .dae import _collada_loaders
16from .obj import _obj_loaders
17from .off import _off_loaders
18from .misc import _misc_loaders
19from .gltf import _gltf_loaders
20from .assimp import _assimp_loaders
21from .threemf import _three_loaders
22from .openctm import _ctm_loaders
23from .xml_based import _xml_loaders
24from .binvox import _binvox_loaders
25from .xyz import _xyz_loaders
26
27
28try:
29    from ..path.exchange.load import load_path, path_formats
30except BaseException as E:
31    # save a traceback to see why path didn't import
32    _path_exception = E
33
34    def load_path(*args, **kwargs):
35        """
36        Dummy load path function that will raise an exception
37        on use. Import of path failed, probably because a
38        dependency is not installed.
39
40        Raises
41        ----------
42        path_exception : BaseException
43          Whatever failed when we imported path
44        """
45        raise _path_exception
46
47    def path_formats():
48        return []
49
50
51def mesh_formats():
52    """
53    Get a list of mesh formats
54
55    Returns
56    -----------
57    loaders : list
58        Extensions of available mesh loaders
59        i.e. 'stl', 'ply', etc.
60    """
61    return list(mesh_loaders.keys())
62
63
64def available_formats():
65    """
66    Get a list of all available loaders
67
68    Returns
69    -----------
70    loaders : list
71        Extensions of available loaders
72        i.e. 'stl', 'ply', 'dxf', etc.
73    """
74    loaders = mesh_formats()
75    loaders.extend(path_formats())
76    loaders.extend(compressed_loaders.keys())
77    return loaders
78
79
80def load(file_obj,
81         file_type=None,
82         resolver=None,
83         **kwargs):
84    """
85    Load a mesh or vectorized path into objects like
86    Trimesh, Path2D, Path3D, Scene
87
88    Parameters
89    -----------
90    file_obj : str, or file- like object
91      The source of the data to be loadeded
92    file_type: str
93      What kind of file type do we have (eg: 'stl')
94    resolver : trimesh.visual.Resolver
95      Object to load referenced assets like materials and textures
96    kwargs : dict
97      Passed to geometry __init__
98
99    Returns
100    ---------
101    geometry : Trimesh, Path2D, Path3D, Scene
102      Loaded geometry as trimesh classes
103    """
104    # check to see if we're trying to load something
105    # that is already a native trimesh Geometry subclass
106    if isinstance(file_obj, Geometry):
107        log.info('Load called on %s object, returning input',
108                 file_obj.__class__.__name__)
109        return file_obj
110
111    # parse the file arguments into clean loadable form
112    (file_obj,  # file- like object
113     file_type,  # str, what kind of file
114     metadata,  # dict, any metadata from file name
115     opened,    # bool, did we open the file ourselves
116     resolver   # object to load referenced resources
117     ) = parse_file_args(file_obj=file_obj,
118                         file_type=file_type,
119                         resolver=resolver)
120
121    try:
122        if isinstance(file_obj, dict):
123            # if we've been passed a dict treat it as kwargs
124            kwargs.update(file_obj)
125            loaded = load_kwargs(kwargs)
126        elif file_type in path_formats():
127            # path formats get loaded with path loader
128            loaded = load_path(file_obj,
129                               file_type=file_type,
130                               **kwargs)
131        elif file_type in mesh_loaders:
132            # mesh loaders use mesh loader
133            loaded = load_mesh(file_obj,
134                               file_type=file_type,
135                               resolver=resolver,
136                               **kwargs)
137        elif file_type in compressed_loaders:
138            # for archives, like ZIP files
139            loaded = load_compressed(file_obj,
140                                     file_type=file_type,
141                                     **kwargs)
142        elif file_type in voxel_loaders:
143            loaded = voxel_loaders[file_type](
144                file_obj,
145                file_type=file_type,
146                resolver=resolver,
147                **kwargs)
148        else:
149            if file_type in ['svg', 'dxf']:
150                # call the dummy function to raise the import error
151                # this prevents the exception from being super opaque
152                load_path()
153            else:
154                raise ValueError('File type: %s not supported' %
155                                 file_type)
156    finally:
157        # close any opened files even if we crashed out
158        if opened:
159            file_obj.close()
160
161    # add load metadata ('file_name') to each loaded geometry
162    for i in util.make_sequence(loaded):
163        i.metadata.update(metadata)
164
165    # if we opened the file in this function ourselves from a
166    # file name clean up after ourselves by closing it
167    if opened:
168        file_obj.close()
169
170    return loaded
171
172
173@log_time
174def load_mesh(file_obj,
175              file_type=None,
176              resolver=None,
177              **kwargs):
178    """
179    Load a mesh file into a Trimesh object
180
181    Parameters
182    -----------
183    file_obj : str or file object
184      File name or file with mesh data
185    file_type : str or None
186      Which file type, e.g. 'stl'
187    kwargs : dict
188      Passed to Trimesh constructor
189
190    Returns
191    ----------
192    mesh : trimesh.Trimesh or trimesh.Scene
193      Loaded geometry data
194    """
195
196    # parse the file arguments into clean loadable form
197    (file_obj,  # file- like object
198     file_type,  # str, what kind of file
199     metadata,  # dict, any metadata from file name
200     opened,    # bool, did we open the file ourselves
201     resolver   # object to load referenced resources
202     ) = parse_file_args(file_obj=file_obj,
203                         file_type=file_type,
204                         resolver=resolver)
205
206    try:
207        # make sure we keep passed kwargs to loader
208        # but also make sure loader keys override passed keys
209        results = mesh_loaders[file_type](file_obj,
210                                          file_type=file_type,
211                                          resolver=resolver,
212                                          **kwargs)
213
214        if util.is_file(file_obj):
215            file_obj.close()
216
217        if not isinstance(results, list):
218            results = [results]
219
220        loaded = []
221        for result in results:
222            kwargs.update(result)
223            loaded.append(load_kwargs(kwargs))
224            loaded[-1].metadata.update(metadata)
225        if len(loaded) == 1:
226            loaded = loaded[0]
227        # show the repr for loaded
228        log.debug('loaded {} using {}'.format(
229            str(loaded),
230            mesh_loaders[file_type].__name__))
231    finally:
232        # if we failed to load close file
233        if opened:
234            file_obj.close()
235
236    return loaded
237
238
239def load_compressed(file_obj,
240                    file_type=None,
241                    resolver=None,
242                    mixed=False,
243                    **kwargs):
244    """
245    Given a compressed archive load all the geometry that
246    we can from it.
247
248    Parameters
249    ----------
250    file_obj : open file-like object
251      Containing compressed data
252    file_type : str
253      Type of the archive file
254    mixed : bool
255      If False, for archives containing both 2D and 3D
256      data will only load the 3D data into the Scene.
257
258    Returns
259    ----------
260    scene : trimesh.Scene
261      Geometry loaded in to a Scene object
262    """
263
264    # parse the file arguments into clean loadable form
265    (file_obj,  # file- like object
266     file_type,  # str, what kind of file
267     metadata,  # dict, any metadata from file name
268     opened,    # bool, did we open the file ourselves
269     resolver   # object to load referenced resources
270     ) = parse_file_args(file_obj=file_obj,
271                         file_type=file_type,
272                         resolver=resolver)
273
274    try:
275        # a dict of 'name' : file-like object
276        files = util.decompress(file_obj=file_obj,
277                                file_type=file_type)
278        # store loaded geometries as a list
279        geometries = []
280
281        # so loaders can access textures/etc
282        resolver = visual.resolvers.ZipResolver(files)
283
284        # try to save the files with meaningful metadata
285        if 'file_path' in metadata:
286            archive_name = metadata['file_path']
287        else:
288            archive_name = 'archive'
289
290        # populate our available formats
291        if mixed:
292            available = available_formats()
293        else:
294            # all types contained in ZIP archive
295            contains = set(util.split_extension(n).lower()
296                           for n in files.keys())
297            # if there are no mesh formats available
298            if contains.isdisjoint(mesh_formats()):
299                available = path_formats()
300            else:
301                available = mesh_formats()
302
303        for name, data in files.items():
304            # only load formats that we support
305            compressed_type = util.split_extension(name).lower()
306            if compressed_type not in available:
307                # don't raise an exception, just try the next one
308                continue
309            # store the file name relative to the archive
310            metadata['file_name'] = (archive_name + '/' +
311                                     os.path.basename(name))
312            # load the individual geometry
313            loaded = load(file_obj=data,
314                          file_type=compressed_type,
315                          resolver=resolver,
316                          metadata=metadata,
317                          **kwargs)
318
319            # some loaders return multiple geometries
320            if util.is_sequence(loaded):
321                # if the loader has returned a list of meshes
322                geometries.extend(loaded)
323            else:
324                # if the loader has returned a single geometry
325                geometries.append(loaded)
326
327    finally:
328        # if we opened the file in this function
329        # clean up after ourselves
330        if opened:
331            file_obj.close()
332
333    # append meshes or scenes into a single Scene object
334    result = append_scenes(geometries)
335
336    return result
337
338
339def load_remote(url, **kwargs):
340    """
341    Load a mesh at a remote URL into a local trimesh object.
342
343    This must be called explicitly rather than automatically
344    from trimesh.load to ensure users don't accidentally make
345    network requests.
346
347    Parameters
348    ------------
349    url : string
350      URL containing mesh file
351    **kwargs : passed to `load`
352    """
353    # import here to keep requirement soft
354    import requests
355
356    # download the mesh
357    response = requests.get(url)
358    # wrap as file object
359    file_obj = util.wrap_as_stream(response.content)
360
361    # so loaders can access textures/etc
362    resolver = visual.resolvers.WebResolver(url)
363
364    # actually load
365    loaded = load(file_obj=file_obj,
366                  file_type=url,
367                  resolver=resolver,
368                  **kwargs)
369    return loaded
370
371
372def load_kwargs(*args, **kwargs):
373    """
374    Load geometry from a properly formatted dict or kwargs
375    """
376    def handle_scene():
377        """
378        Load a scene from our kwargs:
379
380        class:      Scene
381        geometry:   dict, name: Trimesh kwargs
382        graph:      list of dict, kwargs for scene.graph.update
383        base_frame: str, base frame of graph
384        """
385        scene = Scene()
386        scene.geometry.update({k: load_kwargs(v) for
387                               k, v in kwargs['geometry'].items()})
388        for k in kwargs['graph']:
389            if isinstance(k, dict):
390                scene.graph.update(**k)
391            elif util.is_sequence(k) and len(k) == 3:
392                scene.graph.update(k[1], k[0], **k[2])
393        if 'base_frame' in kwargs:
394            scene.graph.base_frame = kwargs['base_frame']
395        if 'metadata' in kwargs:
396            scene.metadata.update(kwargs['metadata'])
397        return scene
398
399    def handle_mesh():
400        """
401        Handle the keyword arguments for a Trimesh object
402        """
403        # if they've been serialized as a dict
404        if (isinstance(kwargs['vertices'], dict) or
405                isinstance(kwargs['faces'], dict)):
406            return Trimesh(**misc.load_dict(kwargs))
407        # otherwise just load that puppy
408        return Trimesh(**kwargs)
409
410    def handle_export():
411        """
412        Handle an exported mesh.
413        """
414        data, file_type = kwargs['data'], kwargs['file_type']
415        if not isinstance(data, dict):
416            data = util.wrap_as_stream(data)
417        k = mesh_loaders[file_type](data,
418                                    file_type=file_type)
419        return Trimesh(**k)
420
421    def handle_pointcloud():
422        return PointCloud(**kwargs)
423
424    # if we've been passed a single dict instead of kwargs
425    # substitute the dict for kwargs
426    if (len(kwargs) == 0 and
427        len(args) == 1 and
428            isinstance(args[0], dict)):
429        kwargs = args[0]
430
431    # (function, tuple of expected keys)
432    # order is important
433    handlers = (
434        (handle_scene, ('graph', 'geometry')),
435        (handle_mesh, ('vertices', 'faces')),
436        (handle_pointcloud, ('vertices',)),
437        (handle_export, ('file_type', 'data')))
438
439    # filter out keys with a value of None
440    kwargs = {k: v for k, v in kwargs.items() if v is not None}
441
442    # loop through handler functions and expected key
443    for func, expected in handlers:
444        if all(i in kwargs for i in expected):
445            # all expected kwargs exist
446            handler = func
447            # exit the loop as we found one
448            break
449    else:
450        raise ValueError('unable to determine type!')
451
452    return handler()
453
454
455def parse_file_args(file_obj,
456                    file_type,
457                    resolver=None,
458                    **kwargs):
459    """
460    Given a file_obj and a file_type try to magically convert
461    arguments to a file-like object and a lowercase string of
462    file type.
463
464    Parameters
465    -----------
466    file_obj : str
467      if string represents a file path, returns:
468        file_obj:   an 'rb' opened file object of the path
469        file_type:  the extension from the file path
470
471     if string is NOT a path, but has JSON-like special characters:
472        file_obj:   the same string passed as file_obj
473        file_type:  set to 'json'
474
475     if string is a valid-looking URL
476        file_obj: an open 'rb' file object with retrieved data
477        file_type: from the extension
478
479     if string is none of those:
480        raise ValueError as we can't do anything with input
481
482     if file like object:
483        ValueError will be raised if file_type is None
484        file_obj:  same as input
485        file_type: same as input
486
487     if other object: like a shapely.geometry.Polygon, etc:
488        file_obj:  same as input
489        file_type: if None initially, set to the class name
490                    (in lower case), otherwise passed through
491
492    file_type : str
493         type of file and handled according to above
494
495    Returns
496    -----------
497    file_obj : file-like object
498      Contains data
499    file_type : str
500      Lower case of the type of file (eg 'stl', 'dae', etc)
501    metadata : dict
502      Any metadata gathered
503    opened : bool
504      Did we open the file or not
505    resolver : trimesh.visual.Resolver
506      Resolver to load other assets
507    """
508    metadata = {}
509    opened = False
510    if ('metadata' in kwargs and
511            isinstance(kwargs['metadata'], dict)):
512        metadata.update(kwargs['metadata'])
513
514    if util.is_pathlib(file_obj):
515        # convert pathlib objects to string
516        file_obj = str(file_obj.absolute())
517
518    if util.is_file(file_obj) and file_type is None:
519        raise ValueError('file_type must be set for file objects!')
520    if util.is_string(file_obj):
521        try:
522            # os.path.isfile will return False incorrectly
523            # if we don't give it an absolute path
524            file_path = os.path.expanduser(file_obj)
525            file_path = os.path.abspath(file_path)
526            exists = os.path.isfile(file_path)
527        except BaseException:
528            exists = False
529
530        # file obj is a string which exists on filesystm
531        if exists:
532            # if not passed create a resolver to find other files
533            if resolver is None:
534                resolver = visual.resolvers.FilePathResolver(file_path)
535            # save the file name and path to metadata
536            metadata['file_path'] = file_path
537            metadata['file_name'] = os.path.basename(file_obj)
538            # if file_obj is a path that exists use extension as file_type
539            if file_type is None:
540                file_type = util.split_extension(
541                    file_path,
542                    special=['tar.gz', 'tar.bz2'])
543            # actually open the file
544            file_obj = open(file_path, 'rb')
545            opened = True
546        else:
547            if '{' in file_obj:
548                # if a dict bracket is in the string, its probably a straight
549                # JSON
550                file_type = 'json'
551            elif 'https://' in file_obj or 'http://' in file_obj:
552                # we've been passed a URL, warn to use explicit function
553                # and don't do network calls via magical pipeline
554                raise ValueError(
555                    'use load_remote to load URL: {}'.format(file_obj))
556            elif file_type is None:
557                raise ValueError('string is not a file: {}'.format(file_obj))
558
559    if file_type is None:
560        file_type = file_obj.__class__.__name__
561
562    if util.is_string(file_type) and '.' in file_type:
563        # if someone has passed the whole filename as the file_type
564        # use the file extension as the file_type
565        if 'file_path' not in metadata:
566            metadata['file_path'] = file_type
567        metadata['file_name'] = os.path.basename(file_type)
568        file_type = util.split_extension(file_type)
569        if resolver is None and os.path.exists(file_type):
570            resolver = visual.resolvers.FilePathResolver(file_type)
571
572    # all our stored extensions reference in lower case
573    file_type = file_type.lower()
574
575    # if we still have no resolver try using file_obj name
576    if (resolver is None and
577        hasattr(file_obj, 'name') and
578        file_obj.name is not None and
579            len(file_obj.name) > 0):
580        resolver = visual.resolvers.FilePathResolver(file_obj.name)
581
582    return file_obj, file_type, metadata, opened, resolver
583
584
585# loader functions for compressed extensions
586compressed_loaders = {'zip': load_compressed,
587                      'tar.bz2': load_compressed,
588                      'tar.gz': load_compressed}
589
590# map file_type to loader function
591mesh_loaders = {}
592# assimp has a lot of loaders, but they are all quite slow
593# load first and replace with native loaders where possible
594mesh_loaders.update(_assimp_loaders)
595mesh_loaders.update(_misc_loaders)
596mesh_loaders.update(_stl_loaders)
597mesh_loaders.update(_ctm_loaders)
598mesh_loaders.update(_ply_loaders)
599mesh_loaders.update(_xml_loaders)
600mesh_loaders.update(_obj_loaders)
601mesh_loaders.update(_off_loaders)
602mesh_loaders.update(_collada_loaders)
603mesh_loaders.update(_gltf_loaders)
604mesh_loaders.update(_three_loaders)
605mesh_loaders.update(_xyz_loaders)
606
607# collect loaders which return voxel types
608voxel_loaders = {}
609voxel_loaders.update(_binvox_loaders)
610