22import os
33import threading
44import time
5+ import xml .etree .ElementTree as ET
56import zipfile
67from collections .abc import Callable
78from concurrent .futures import ThreadPoolExecutor
2223# Local application imports
2324from . import geode_functions
2425from .geode_objects import geode_objects
26+ from .geode_objects .geode_model import GeodeModel
2527from .geode_objects .geode_object import GeodeObject
2628
2729
@@ -191,11 +193,69 @@ def create_data_folder_from_id(data_id: str) -> str:
191193 return data_path
192194
193195
196+ def model_components (
197+ data_id : str , model : GeodeModel , viewable_file : str
198+ ) -> dict [str , Any ]:
199+ vtm_file_path = geode_functions .data_file_path (data_id , viewable_file )
200+ tree = ET .parse (vtm_file_path )
201+ root = tree .find ("vtkMultiBlockDataSet" )
202+ if root is None :
203+ flask .abort (500 , "Failed to read viewable file" )
204+ uuid_to_flat_index = {}
205+ current_index = 0
206+ assert root is not None
207+ for elem in root .iter ():
208+ if "uuid" in elem .attrib and elem .tag == "DataSet" :
209+ uuid_to_flat_index [elem .attrib ["uuid" ]] = current_index
210+ current_index += 1
211+ model_mesh_components = model .mesh_components ()
212+ mesh_components = []
213+ for mesh_component , ids in model_mesh_components .items ():
214+ component_type = mesh_component .get ()
215+ for id in ids :
216+ geode_id = id .string ()
217+ component_name = geode_id
218+ viewer_id = uuid_to_flat_index [geode_id ]
219+ boundaries = model .boundaries (id )
220+ boundaries_uuid = [boundary .id ().string () for boundary in boundaries ]
221+ internals = model .internals (id )
222+ internals_uuid = [internal .id ().string () for internal in internals ]
223+ mesh_component_object = {
224+ "viewer_id" : viewer_id ,
225+ "geode_id" : geode_id ,
226+ "name" : component_name ,
227+ "type" : component_type ,
228+ "boundaries" : boundaries_uuid ,
229+ "internals" : internals_uuid ,
230+ }
231+ mesh_components .append (mesh_component_object )
232+
233+ model_collection_components = model .collection_components ()
234+ collection_components = []
235+ for collection_component , ids in model_collection_components .items ():
236+ component_type = collection_component .get ()
237+ for id in ids :
238+ geode_id = id .string ()
239+ items = model .items (id )
240+ items_uuid = [item .id ().string () for item in items ]
241+ collection_component_object = {
242+ "geode_id" : geode_id ,
243+ "name" : geode_id ,
244+ "type" : component_type ,
245+ "items" : items_uuid ,
246+ }
247+ collection_components .append (collection_component_object )
248+ return {
249+ "mesh_components" : mesh_components ,
250+ "collection_components" : collection_components ,
251+ }
252+
253+
194254def save_all_viewables_and_return_info (
195255 geode_object : GeodeObject ,
196256 data : Data ,
197257 data_path : str ,
198- ) -> dict [str , str | list [ str ] ]:
258+ ) -> dict [str , Any ]:
199259 with ThreadPoolExecutor () as executor :
200260 native_files , viewable_path , light_path = executor .map (
201261 lambda args : args [0 ](args [1 ]),
@@ -225,7 +285,8 @@ def save_all_viewables_and_return_info(
225285 name = geode_object .identifier .name ()
226286 if not name :
227287 flask .abort (400 , "Geode object has no name defined." )
228- return {
288+
289+ response : dict [str , Any ] = {
229290 "native_file" : data .native_file ,
230291 "viewable_file" : data .viewable_file ,
231292 "id" : data .id ,
@@ -234,11 +295,14 @@ def save_all_viewables_and_return_info(
234295 "binary_light_viewable" : binary_light_viewable .decode ("utf-8" ),
235296 "geode_object_type" : data .geode_object ,
236297 }
298+ if isinstance (geode_object , GeodeModel ):
299+ response |= model_components (data .id , geode_object , data .viewable_file )
300+ return response
237301
238302
239303def generate_native_viewable_and_light_viewable_from_object (
240304 geode_object : GeodeObject ,
241- ) -> dict [str , str | list [ str ] ]:
305+ ) -> dict [str , Any ]:
242306 data = Data .create (
243307 geode_object = geode_object .geode_object_type (),
244308 viewer_object = geode_object .viewer_type (),
@@ -250,7 +314,7 @@ def generate_native_viewable_and_light_viewable_from_object(
250314
251315def generate_native_viewable_and_light_viewable_from_file (
252316 geode_object_type : GeodeObjectType , input_file : str
253- ) -> dict [str , str | list [ str ] ]:
317+ ) -> dict [str , Any ]:
254318 generic_geode_object = geode_objects [geode_object_type ]
255319 data = Data .create (
256320 geode_object = geode_object_type ,
0 commit comments