extract_maya_usd

class ExtractMayaUsd(plugin.MayaExtractorPlugin,
                     publish.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"]

    # 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,
        }

    @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": "catmullClark",
            "stripNamespaces": True,
            "mergeTransformAndShape": True,
            "exportDisplayColor": False,
            "exportColorSets": True,
            "exportInstances": True,
            "exportUVs": True,
            "exportVisibility": True,
            "exportComponentTags": False,
            "exportRefsAsInstanceable": False,
            "eulerFilter": True,
            "renderableOnly": False,
            "shadingMode": "none",
            "convertMaterialsTo": "none",
            "jobContext": None,
            "filterTypes": None,
            "staticSingleSample": True,
            "worldspace": True
        }

    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)

        attr_defs = [
            UISeparatorDef("sep_usd_options"),
            UILabelDef("USD Options"),
        ]
        attr_defs.extend(
            super().get_attr_defs_for_instance(create_context, instance)
        )
        attr_defs.extend(cls._get_additional_attr_defs(is_enabled))
        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) -> list:
        return [
            BoolDef("stripNamespaces",
                    label="Strip Namespaces",
                    tooltip="Strip Namespaces in the USD Export",
                    visible=visible,
                    default=True),
            BoolDef("worldspace",
                    label="World-Space",
                    tooltip="Export all root prim using their full worldspace "
                            "transform instead of their local transform.",
                    visible=visible,
                    default=True),
            BoolDef("exportComponentTags",
                    label="Export Component Tags",
                    tooltip="When enabled, export any geometry component tags "
                            "as UsdGeomSubset data.",
                    visible=visible,
                    default=False),
            BoolDef("exportVisibility",
                    label="Export Visibility",
                    tooltip="Export any state and animation on Maya visibility"
                            " attributes.",
                    visible=visible,
                    default=True),
            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=True),
            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="catmullClark"
            )
        ]

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