Jedd Morgan
141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
from typing import Dict, Optional, Union
|
|
|
|
import bpy
|
|
from attrs import define
|
|
from bpy.types import ID, Collection, Object
|
|
from specklepy.objects.base import Base
|
|
from specklepy.objects.graph_traversal.commit_object_builder import (
|
|
ROOT,
|
|
CommitObjectBuilder,
|
|
)
|
|
from specklepy.objects.other import Collection as SCollection
|
|
|
|
from bpy_speckle.functions import _report
|
|
|
|
ELEMENTS = "elements"
|
|
|
|
|
|
def _id(native_object: ID) -> str:
|
|
# NOTE: to avoid naming collisions, we prefix collections and objects differently
|
|
return f"{type(native_object).__name__}:{native_object.name_full}"
|
|
|
|
|
|
def _try_id(native_object: Optional[Union[Collection, Object]]) -> Optional[str]:
|
|
return _id(native_object) if native_object else None
|
|
|
|
|
|
def convert_collection_to_speckle(col: Collection) -> SCollection:
|
|
converted_collection = SCollection(
|
|
name=col.name_full, collectionType="Blender Collection", elements=[]
|
|
)
|
|
converted_collection.applicationId = _id(col)
|
|
|
|
color_tag = col.color_tag
|
|
if color_tag and color_tag != "NONE":
|
|
converted_collection["colorTag"] = col.color_tag
|
|
|
|
return converted_collection
|
|
|
|
|
|
@define(slots=True)
|
|
class BlenderCommitObjectBuilder(CommitObjectBuilder[Object]):
|
|
_collections: Dict[str, SCollection]
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._collections = {}
|
|
|
|
def include_object(self, conversion_result: Base, native_object: Object) -> None:
|
|
# Set the Child -> Parent relationships
|
|
parent = native_object.parent
|
|
|
|
parent_collections = native_object.users_collection
|
|
parent_collection = (
|
|
parent_collections[0] if len(parent_collections) > 0 else None
|
|
) # NOTE: we don't support objects appearing in more than one collection, for now, we will just take the zeroth one
|
|
|
|
app_id = _id(native_object)
|
|
conversion_result.applicationId = app_id
|
|
self.converted[app_id] = conversion_result
|
|
|
|
# in order or priority, direct parent, direct parent collection, root
|
|
self.set_relationship(
|
|
app_id,
|
|
(_try_id(parent), ELEMENTS),
|
|
(_try_id(parent_collection), ELEMENTS),
|
|
(ROOT, ELEMENTS),
|
|
)
|
|
# if parent_collection:
|
|
# self._include_collection(parent_collection)
|
|
|
|
def ensure_collection(self, col: Collection) -> SCollection:
|
|
id = _id(col)
|
|
if id in self._collections:
|
|
return self._collections[id] # collection already converted!
|
|
|
|
# Set the Parent -> Children relationships
|
|
for c in col.children:
|
|
# NOTE: There's no falling back to the grandparent, if the direct parent collection wasn't converted, then we we fallback to the root
|
|
self.set_relationship(_id(c), (id, ELEMENTS), (ROOT, ELEMENTS))
|
|
|
|
# Set Child -> Parent relationship
|
|
# parent = self.find_collection_parent(col)
|
|
# self.set_relationship(id, (_try_builder_id(parent), ELEMENTS), (ROOT, ELEMENTS))
|
|
|
|
converted_collection = convert_collection_to_speckle(col)
|
|
self.converted[id] = converted_collection
|
|
self._collections[id] = converted_collection
|
|
|
|
return converted_collection
|
|
|
|
def build_commit_object(self, root_commit_object: Base) -> None:
|
|
assert root_commit_object.applicationId in self.converted
|
|
|
|
# Create all collections
|
|
root_col = self.ensure_collection(bpy.context.scene.collection)
|
|
root_col.collectionType = "Scene Collection"
|
|
for col in bpy.context.scene.collection.children_recursive:
|
|
self.ensure_collection(col)
|
|
|
|
objects_to_build = set(self.converted.values())
|
|
objects_to_build.remove(root_commit_object)
|
|
|
|
self.apply_relationships(objects_to_build, root_commit_object)
|
|
|
|
assert isinstance(root_commit_object, SCollection)
|
|
# Kill unused collections
|
|
|
|
def should_remove_unuseful_collection(
|
|
col: SCollection,
|
|
) -> bool: # TODO: this maybe could be optimised
|
|
elements = col.elements
|
|
if not elements:
|
|
return True
|
|
|
|
should_remove_this_col = True
|
|
|
|
i = 0
|
|
while i < len(elements):
|
|
c = elements[i]
|
|
if not isinstance(c, SCollection):
|
|
# col has objects (c)
|
|
should_remove_this_col = False
|
|
i += 1
|
|
continue
|
|
|
|
if should_remove_unuseful_collection(c):
|
|
# c is not useful, kill it
|
|
del elements[i]
|
|
else:
|
|
# col has a child (c) with objects
|
|
should_remove_this_col = False
|
|
i += 1
|
|
continue
|
|
|
|
return should_remove_this_col
|
|
|
|
if should_remove_unuseful_collection(root_commit_object):
|
|
_report(
|
|
"WARNING: Only empty collections have been converted!"
|
|
) # TODO: consider raising exception here, to halt the send operation
|