extract_maya_usd

class ExtractMayaUsd(plugin.MayaExtractorPlugin,
                     OptionalPyblishPluginMixin):
    """Extractor for Maya USD Asset data.

    Upon publish a .usd (or .usdz) asset file will typically be written.
    """

    enabled = True
    label = "Extract Maya USD Asset"
    families = ["mayaUsd"]

    # Settings
    overrides = []
    stripNamespaces: bool = True
    worldspace: bool = True
    exportComponentTags: bool = False
    exportVisibility: bool = True
    mergeTransformAndShape: bool = True
    defaultMeshScheme: str = "catmullClark"
    exportBlendShapes: bool = True
    exportSkin: str = "none"
    exportSkels: str = "none"
    exportCollectionBasedBindings: bool = False
    exportColorSets: bool = True
    exportDisplayColor: bool = False
    exportMaterials: bool = True
    exportAssignedMaterials: bool = False
    exportInstances: bool = True
    exportUVs: bool = True
    upAxis: str = "mayaPrefs"
    unit: str = "mayaPrefs"

    # Default prefix for custom attributes to USD attributes
    # if no other mapping is provided
    custom_attr_namespace: str = ""
    # Direct attribute to USD attribute name mapping
    # if attribute name not specified in custom_attr_mapping
    custom_attr_name_mapping: list[dict[str, str]]
    # Explicit attribute mapping matching the Maya USD Export
    # USD_UserExportedAttributesJson data structure
    custom_attr_mapping: str

    @property
    def options(self):
        """Overridable options for Maya USD Export

        Given in the following format
            - {NAME: EXPECTED TYPE}

        If the overridden option's type does not match,
        the option is not included and a warning is logged.

        """

        # TODO: Support more `mayaUSDExport` parameters
        return {
            "chaser": (list, None),  # optional list
            "chaserArgs": (list, None),  # optional list
            "defaultUSDFormat": str,
            "defaultMeshScheme": str,
            "stripNamespaces": bool,
            "mergeTransformAndShape": bool,
            "exportDisplayColor": bool,
            "exportColorSets": bool,
            "exportInstances": bool,
            "exportUVs": bool,
            "exportVisibility": bool,
            "exportComponentTags": bool,
            "exportRefsAsInstanceable": bool,
            "eulerFilter": bool,
            "renderableOnly": bool,
            "convertMaterialsTo": str,
            "shadingMode": (str, None),  # optional str
            "jobContext": (list, None),  # optional list
            "filterTypes": (list, None),  # optional list
            "staticSingleSample": bool,
            "worldspace": bool,
            "exportCollectionBasedBindings": bool,
            "exportBlendShapes": bool,
            "exportSkels": str,
            "exportSkin": str,
            "exportMaterials": bool,
            "exportAssignedMaterials": bool,
            "upAxis": str,
            "unit": (str, None),  # optional str
        }

    @property
    def default_options(self):
        """The default options for Maya USD Export."""

        # TODO: Support more `mayaUSDExport` parameters
        return {
            "chaser": None,
            "chaserArgs": None,
            "defaultUSDFormat": "usdc",
            "defaultMeshScheme": self.defaultMeshScheme,
            "stripNamespaces": self.stripNamespaces,
            "mergeTransformAndShape": self.mergeTransformAndShape,
            "exportDisplayColor": self.exportDisplayColor,
            "exportColorSets": self.exportColorSets,
            "exportInstances": self.exportInstances,
            "exportUVs": self.exportUVs,
            "exportVisibility": self.exportVisibility,
            "exportComponentTags": self.exportComponentTags,
            "exportRefsAsInstanceable": False,
            "eulerFilter": True,
            "renderableOnly": False,
            "shadingMode": "none",
            "convertMaterialsTo": "none",
            "jobContext": None,
            "filterTypes": None,
            "staticSingleSample": True,
            "worldspace": self.worldspace,
            "exportCollectionBasedBindings": (
                self.exportCollectionBasedBindings
            ),
            "exportBlendShapes": self.exportBlendShapes,
            "exportSkels": self.exportSkels,
            "exportSkin": self.exportSkin,
            "exportMaterials": self.exportMaterials,
            "exportAssignedMaterials": self.exportAssignedMaterials,
            "upAxis": self.upAxis,
            "unit": self.unit,
        }

    def parse_overrides(self, overrides, options):
        """Inspect data of instance to determine overridden options"""

        for key in overrides:
            if key not in self.options:
                continue

            # Ensure the data is of correct type
            value = overrides[key]
            if isinstance(value, str):
                value = str(value)
            if not isinstance(value, self.options[key]):
                self.log.warning(
                    "Overridden attribute {key} was of "
                    "the wrong type: {invalid_type} "
                    "- should have been {valid_type}".format(
                        key=key,
                        invalid_type=type(value).__name__,
                        valid_type=self.options[key].__name__))
                continue

            options[key] = value

        # Do not pass None values
        for key, value in options.copy().items():
            if value is None:
                del options[key]

        return options

    def filter_members(self, members):
        # Can be overridden by inherited classes
        return members

    def process(self, instance):
        if not self.is_active(instance.data):
            return

        if instance.data.get("farm"):
            self.log.debug("Should be processed on farm, skipping.")
            return

        attr_values = self.get_attr_values_from_data(instance.data)

        # Load plugin first
        cmds.loadPlugin("mayaUsdPlugin", quiet=True)

        # Define output file path
        staging_dir = self.staging_dir(instance)
        file_name = "{0}.usd".format(instance.name)
        file_path = os.path.join(staging_dir, file_name)
        file_path = file_path.replace('\\', '/')

        # Parse export options
        options = self.default_options
        options = self.parse_overrides(instance.data, options)
        options = self.parse_overrides(attr_values, options)

        # Perform extraction
        self.log.debug("Performing extraction ...")

        members = instance.data("setMembers")
        self.log.debug('Collected objects: {}'.format(members))
        members = self.filter_members(members)
        if not members:
            self.log.error('No members!')
            return

        export_anim_data = instance.data.get("exportAnimationData", True)
        start = instance.data.get("frameStartHandle", 0)

        if export_anim_data:
            end = instance.data["frameEndHandle"]
            options["frameRange"] = (start, end)
            options["frameStride"] = instance.data.get("step", 1.0)

        if instance.data.get("exportRoots", True):
            # Do not include 'objectSets' as roots because the export command
            # will fail. We only include the transforms among the members.
            options["exportRoots"] = cmds.ls(members,
                                             type="transform",
                                             long=True)
        else:
            options["selection"] = True

        # TODO: Remove hardcoded filterTypes
        # We always filter constraint types because they serve no valuable
        # data (it doesn't preserve the actual constraint) but it does
        # introduce the problem that Shapes do not merge into the Transform
        # on export anymore because they are usually parented under transforms
        # See: https://github.com/Autodesk/maya-usd/issues/2070
        options["filterTypes"] = ["constraint"]

        def parse_attr_str(attr_str):
            """Return list of strings from `a,b,c,d` to `[a, b, c, d]`.

            Args:
                attr_str (str): Concatenated attributes by comma

            Returns:
                List[str]: list of attributes
            """
            result = list()
            for attr in attr_str.split(","):
                attr = attr.strip()
                if not attr:
                    continue
                result.append(attr)
            return result

        attrs = parse_attr_str(instance.data.get("attr", ""))
        attrs += instance.data.get("userDefinedAttributes", [])
        attrs += ["cbId"]
        attr_prefixes = parse_attr_str(instance.data.get("attrPrefix", ""))

        # Remove arguments for Maya USD versions not supporting them yet
        # Note: Maya 2022.3 ships with Maya USD 0.13.0.
        # TODO: Remove this backwards compatibility if Maya 2022 support is
        #   dropped
        maya_usd_version = parse_version(
            cmds.pluginInfo("mayaUsdPlugin", query=True, version=True)
        )
        for key, required_minimal_version in {
            "exportComponentTags": (0, 14, 0),
            "jobContext": (0, 15, 0),
            "worldspace": (0, 21, 0)
        }.items():
            if key in options and maya_usd_version < required_minimal_version:
                self.log.warning(
                    "Ignoring export flag '%s' because Maya USD version "
                    "%s is lower than minimal supported version %s.",
                    key,
                    maya_usd_version,
                    required_minimal_version
                )
                del options[key]

        # Fix default prim bug in Maya USD 0.30.0 where prefixed `|` remains
        # See: https://github.com/Autodesk/maya-usd/issues/3991
        if (
                options.get("exportRoots")          # only if roots are defined
                and "defaultPrim" not in options    # ignore if already set
                and "rootPrim" not in options       # ignore if root is created
                and maya_usd_version == (0, 30, 0)  # only for Maya USD 0.30.0
        ):
            # Define the default prim name as it will end up in the USD file
            # from the first export root node
            first_root = options["exportRoots"][0]
            default_prim = first_root.rsplit("|", 1)[-1]
            if options["stripNamespaces"]:
                default_prim = default_prim.rsplit(":", 1)[-1]
            options["defaultPrim"] = default_prim

        # Specify custom attribute mapping from settings
        custom_attr_mapping: dict[str, dict[str, Any]] = {}
        try:
            custom_attr_mapping = json.loads(self.custom_attr_mapping)
        except json.decoder.JSONDecodeError:
            pass

        for data in self.custom_attr_name_mapping:
            maya_name: str = data["name"]
            if maya_name in custom_attr_mapping:
                continue
            custom_attr_mapping[maya_name] = {"usdAttrName": data["usd_name"]}

        self.log.debug(
            f"Custom attribute mapping: {custom_attr_mapping}"
        )
        self.log.debug(
            f"Custom attribute default namespace: {self.custom_attr_namespace}"
        )

        # Remove attributes from custom mapping which we do not intend
        # to include in the export
        attrs_lookup = set(attrs)
        for attr_name in list(custom_attr_mapping):
            # Exclude any keys not matching specified `attrs` or
            # `attr_prefixes` because we only want to include them in the
            # export if these are marked as attributes to export. Maya USD
            # exports everything in the custom mapping, so we pop them.
            if attr_name in attrs_lookup:
                continue
            if attr_name.startswith(tuple(attr_prefixes)):
                continue
            del custom_attr_mapping[attr_name]

        self.log.debug("Export options: {0}".format(options))
        self.log.debug('Exporting USD: {} / {}'.format(file_path, members))
        with maintained_time():
            with maintained_selection():
                if not export_anim_data:
                    # Use start frame as current time
                    cmds.currentTime(start)

                with usd_export_attributes(
                    instance[:],
                    attrs=attrs,
                    attr_prefixes=attr_prefixes,
                    custom_attr_default_namespace=self.custom_attr_namespace,
                    mapping=custom_attr_mapping
                ):
                    cmds.select(members, replace=True, noExpand=True)
                    cmds.mayaUSDExport(file=file_path,
                                       **options)

        representation = {
            'name': "usd",
            'ext': "usd",
            'files': file_name,
            'stagingDir': staging_dir
        }
        instance.data.setdefault("representations", []).append(representation)

        self.log.debug(
            "Extracted instance {} to {}".format(instance.name, file_path)
        )

    @classmethod
    def register_create_context_callbacks(cls, create_context):
        create_context.add_value_changed_callback(cls.on_values_changed)

    @classmethod
    def on_values_changed(cls, event):
        """Update instance attribute definitions on attribute changes."""
        for instance_change in event["changes"]:
            # First check if there's a change we want to respond to
            instance = instance_change["instance"]
            if instance is None:
                # Change is on context
                continue

            # Check if active state is toggled
            value_changes = instance_change["changes"]
            if "publish_attributes" not in value_changes:
                continue

            publish_attributes = value_changes["publish_attributes"]
            class_name = cls.__name__
            if class_name not in publish_attributes:
                continue

            if "active" not in publish_attributes[class_name]:
                continue

            # Update the attribute definitions
            new_attrs = cls.get_attr_defs_for_instance(
                event["create_context"], instance
            )
            instance.set_publish_plugin_attr_defs(class_name, new_attrs)

    @classmethod
    def get_attr_defs_for_instance(cls, create_context, instance):
        is_enabled = cls.enabled
        if not is_enabled:
            return []

        if not cls.instance_matches_plugin_families(instance):
            return []

        if cls.optional:
            plugin_attr_values = (
                instance.data
                .get("publish_attributes", {})
                .get(cls.__name__, {})
            )
            is_enabled = plugin_attr_values.get("active", cls.active)
        defs = super().get_attr_defs_for_instance(create_context, instance)
        if not cls.overrides:
            return defs

        attr_defs = [
            UISeparatorDef("sep_usd_options"),
            UILabelDef("USD Options"),
        ]
        attr_defs.extend(defs)
        # The Arguments that can be modified by the Publisher
        override_defs = cls._get_additional_attr_defs(is_enabled)
        overrides = set(cls.overrides)
        for key, value in override_defs.items():
            if key not in overrides:
                continue

            attr_defs.append(value)

        attr_defs.append(
            UISeparatorDef("sep_usd_options_end")
        )
        return attr_defs

    @classmethod
    def convert_attribute_values(cls, create_context, instance):
        # Convert creator attribute 'mergeTransformAndShape' to
        # plugin attribute, because this attribute has moved from
        # the `io.openpype.creators.maya.mayausd` creator to this extractor
        super().convert_attribute_values(create_context, instance)
        if (
                not cls.enabled
                or not instance
                or not cls.instance_matches_plugin_families(instance)
        ):
            return
        if (
                instance.data.get("creator_identifier")
                != "io.openpype.creators.maya.mayausd"
        ):
            return
        creator_attributes = instance.data.get("creator_attributes", {})
        if not creator_attributes:
            return

        keys = ["mergeTransformAndShape"]
        for key in keys:
            if key in creator_attributes:
                # Set attribute value for this plugin
                value = creator_attributes.pop(key)
                class_name = cls.__name__
                instance.publish_attributes[class_name][key] = value

    @classmethod
    def _get_additional_attr_defs(cls, visible: bool) -> dict:
        return {
            "stripNamespaces": BoolDef(
                "stripNamespaces",
                label="Strip Namespaces",
                tooltip="Strip Namespaces in the USD Export",
                visible=visible,
                default=cls.stripNamespaces,
            ),
            "worldspace": BoolDef(
                "worldspace",
                label="World-Space",
                tooltip="Export all root prim using their full worldspace "
                        "transform instead of their local transform.",
                visible=visible,
                default=cls.worldspace,
            ),
            "exportComponentTags": BoolDef(
                "exportComponentTags",
                label="Export Component Tags",
                tooltip="When enabled, export any geometry component tags "
                        "as UsdGeomSubset data.",
                visible=visible,
                default=cls.exportComponentTags,
            ),
            "exportVisibility": BoolDef(
                "exportVisibility",
                label="Export Visibility",
                tooltip="Export any state and animation on Maya visibility"
                        " attributes.",
                visible=visible,
                default=cls.exportVisibility,
            ),
            "mergeTransformAndShape": BoolDef(
                "mergeTransformAndShape",
                label="Merge Transform and Shape",
                tooltip=(
                    "Combine Maya transform and shape into a single USD"
                    "prim that has transform and geometry, for all"
                    " \"geometric primitives\" (gprims).\n"
                    "This results in smaller and faster scenes. Gprims "
                    "will be \"unpacked\" back into transform and shape "
                    "nodes when imported into Maya from USD."
                ),
                visible=visible,
                default=cls.mergeTransformAndShape,
            ),
            "defaultMeshScheme": EnumDef(
                "defaultMeshScheme",
                label="Default Subdivision Method",
                items=[
                    {"value": "catmullClark", "label": "Catmull Clark"},
                    {"value": "loop", "label": "Loop"},
                    {"value": "bilinear", "label": "Bilinear"},
                    {"value": "none", "label": "None"},
                ],
                tooltip=(
                    "Default subdivision method for meshes.\n"
                    "Options are: catmullClark, loop, bilinear, none."
                    "\n\n"
                    "To specify per mesh subdivision schemes add a "
                    "USD_ATTR_subdivisionScheme attribute."
                ),
                visible=visible,
                default=cls.defaultMeshScheme,
            ),
            "exportBlendShapes": BoolDef(
                "exportBlendShapes",
                label="Export Blendshapes",
                tooltip="Enable or disable export of blend shapes.",
                visible=visible,
                default=cls.exportBlendShapes,
            ),
            "exportSkin": EnumDef("exportSkin",
                    label="Export Skin",
                    items=[
                        {"value": "none", "label": "None"},
                        {"value": "auto", "label": "Auto"},
                        {"value": "explicit", "label": "Explicit"},
                    ],
                    tooltip=(
                        "Export skin cluster data.\n\n"
                        "None: do not export skin clusters.\n"
                        "Auto: export skin clusters if present.\n"
                        "Explicit: only export explicitly tagged skin clusters."
                    ),
                    visible=visible,
                    default=cls.exportSkin,
            ),
            "exportSkels": EnumDef("exportSkels",
                    label="Export Skeletons",
                    items=[
                        {"value": "none", "label": "None"},
                        {"value": "auto", "label": "Auto"},
                        {"value": "explicit", "label": "Explicit"},
                    ],
                    tooltip=(
                        "Determines how to export skeletons.\n\n"
                        "None: No skeleton are exported.\n"
                        "Auto: All skeletons will be exported, SkelRoots"
                        " may be created.\n"
                        "Explicit: only export those under SkelRoots."
                    ),
                    visible=visible,
                    default=cls.exportSkels,
            ),
            "exportCollectionBasedBindings": BoolDef(
                "exportCollectionBasedBindings",
                label="Export Collection Based Bindings",
                tooltip=(
                    "Enable or disable export of collection-based material "
                    "assigments. If this option is enabled, export of material "
                    "collections (-mcs) is also enabled, which causes collections "
                    "representing sets of geometry with the same material binding "
                    "to be exported. Materials are bound to the created collections "
                    "on the prim at materialCollectionsPath (specfied via the -mcp "
                    "option). Direct (or per-gprim) bindings are not authored when "
                    "collection-based bindings are enabled."
                ),
                visible=visible,
                default=cls.exportCollectionBasedBindings,
            ),
            "exportColorSets": BoolDef(
                "exportColorSets",
                label="Export Color Sets",
                tooltip="Enable or disable the export of color sets.",
                visible=visible,
                default=cls.exportColorSets,
            ),
            "exportDisplayColor": BoolDef(
                "exportDisplayColor",
                label="Export Display Color",
                tooltip="Enable or disable the export of display color.",
                visible=visible,
                default=cls.exportDisplayColor,
            ),
            "exportMaterials": BoolDef(
                "exportMaterials",
                label="Export Materials",
                tooltip="Enable or disable the export of materials.",
                visible=visible,
                default=cls.exportMaterials,
            ),
            "exportAssignedMaterials": BoolDef(
                "exportAssignedMaterials",
                label="Export Assigned Materials",
                tooltip="Export materials only if they are assigned to a mesh.",
                visible=visible,
                default=cls.exportAssignedMaterials,
            ),
            "exportInstances": BoolDef(
                "exportInstances",
                label="Export Instances",
                tooltip="Enable or disable the export of instances.",
                visible=visible,
                default=cls.exportInstances,
            ),
            "exportUVs": BoolDef(
                "exportUVs",
                label="Export UVs",
                tooltip="Enable or disable the export of UV sets.",
                visible=visible,
                default=cls.exportUVs,
            ),
            "upAxis": EnumDef(
                "upAxis",
                label="Up Axis",
                items=[
                    {"value": "mayaPrefs", "label": "Maya Preferences"},
                    {"value": "none", "label": "None"},
                    {"value": "y", "label": "Y"},
                    {"value": "z", "label": "Z"},
                ],
                tooltip=(
                    "How the up-axis of the exported USD is controlled. "
                    "\"mayaPrefs\" follows the current Maya Preferences. "
                    "\"none\" does not author up-axis. \"y\" or \"z\" author "
                    "that axis and convert data if the Maya preferences does "
                    "not match."
                ),
                visible=visible,
                default=cls.upAxis,
            ),
            "unit": EnumDef(
                "unit",
                label="Export Unit",
                items=[
                    {"value": "mayaPrefs", "label": "Maya Preferences"},
                    {"value": "none", "label": "None"},
                    {"value": "mm", "label": "Millimeters"},
                    {"value": "cm", "label": "Centimeters"},
                    {"value": "m", "label": "Meters"},
                    {"value": "in", "label": "Inches"},
                    {"value": "ft", "label": "Feet"},
                ],
                tooltip=(
                    "How the measuring units of the exported USD is controlled. "
                    "\"mayaPrefs\" follows the current Maya Preferences. "
                    "\"none\" does not author units. Explicit units (cm, inch, etc) "
                    "author that and convert data if the Maya preferences does "
                    "not match."
                ),
                visible=visible,
                default=cls.unit,
            )
        }

class ExtractMayaUsdAnim(ExtractMayaUsd):
    """Extractor for Maya USD Animation Sparse Cache data.

    This will extract the sparse cache data from the scene and generate a
    USD file with all the animation data.

    Upon publish a .usd sparse cache will be written.
    """
    label = "Extract USD Animation"
    families = ["animation"]

    # Exposed in settings
    optional = True
    active = False

    # TODO: Support writing out point deformation only, avoid writing UV sets
    #       component tags and potentially remove `faceVertexCounts`,
    #       `faceVertexIndices` and `doubleSided` parameters as well.
    def filter_members(self, members):
        out_set = next((i for i in members if i.endswith("out_SET")), None)

        if out_set is None:
            self.log.warning("Expecting out_SET")
            return None

        members = cmds.ls(cmds.sets(out_set, query=True), long=True)
        return members

class ExtractMayaUsdModel(ExtractMayaUsd):
    """Extractor for Maya USD Asset data for model family

    Upon publish a .usd (or .usdz) asset file will typically be written.
    """

    label = "Extract USD"
    families = ["model"]

    # Exposed in settings
    optional = True
    active = False

    def process(self, instance):
        # TODO: Fix this without changing instance data
        instance.data["exportAnimationData"] = False
        super().process(instance)

class ExtractMayaUsdPointcache(ExtractMayaUsd):
    """Extractor for Maya USD for 'pointcache' family"""

    label = "Extract USD"
    families = ["pointcache"]

    # Exposed in settings
    optional = True
    active = False