From 94b220bb9cd490b0c4cdafe7adab8d46c4f3edf5 Mon Sep 17 00:00:00 2001
From: Marcus Ottosson <konstruktion@gmail.com>
Date: Sat, 10 Apr 2021 08:26:11 +0100
Subject: [PATCH] More protection, support nice name of static attributes

---
 cmdx.py | 44 +++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 5 deletions(-)

diff --git a/cmdx.py b/cmdx.py
index 5cf84b7..0750e88 100644
--- a/cmdx.py
+++ b/cmdx.py
@@ -1653,6 +1653,7 @@ def parent(self, type=None):
         if not type or type == self._fn.__class__(mobject).typeName:
             return cls(mobject)
 
+    @protected
     def lineage(self, type=None):
         """Yield parents all the way up a hierarchy
 
@@ -1677,6 +1678,7 @@ def lineage(self, type=None):
             yield parent
             parent = parent.parent(type)
 
+    @protected
     def children(self,
                  type=None,
                  filter=om.MFn.kTransform,
@@ -3111,20 +3113,51 @@ def niceName(self):
         """The nice name of this plug, visible in e.g. Channel Box
 
         Examples:
-            >>> node = createNode("transform")
+            >>> _new()
+            >>> node = createNode("transform", name="myTransform")
+
+            # Return pairs of nice names for compound attributes
+            >>> node["scale"].niceName == ("Scale X", "Scale Y", "Scale Z")
+            True
+
             >>> assert node["translateY"].niceName == "Translate Y"
             >>> node["translateY"].niceName = "New Name"
             >>> assert node["translateY"].niceName == "New Name"
 
+            # The nice name is preserved on scene open
+            >>> _save()
+            >>> _load()
+            >>> node = encode("myTransform")
+            >>> assert node["translateY"].niceName == "New Name"
+
         """
 
         # No way of retrieving this information via the API?
-        return cmds.attributeName(self.path(), nice=True)
+
+        if self.isArray or self.isCompound:
+            return tuple(
+                cmds.attributeName(plug.path(), nice=True)
+                for plug in self
+            )
+        else:
+            return cmds.attributeName(self.path(), nice=True)
 
     @niceName.setter
     def niceName(self, value):
-        fn = om.MFnAttribute(self._mplug.attribute())
-        fn.setNiceNameOverride(value)
+        elements = (
+            self
+            if self.isArray or self.isCompound
+            else [self]
+        )
+
+        for el in elements:
+            if el._mplug.isDynamic:
+                # Use setAttr as isKeyable doesn't
+                # persist on scene save for dynamic attributes.
+                cmds.addAttr(el.path(), edit=True, niceName=value)
+            else:
+                fn = om.MFnAttribute(el._mplug.attribute())
+                fn.setNiceNameOverride(value)
 
     @property
     def default(self):
@@ -5759,6 +5792,7 @@ class DagModifier(_BaseModifier):
         >>> mod = DagModifier()
         >>> parent = mod.createNode("transform", name="myParent")
         >>> child = mod.createNode("transform", name="myChild", parent=parent)
+        >>> _ = mod.createNode("transform", name="keepAlive", parent=parent)
         >>> mod.doIt()
         >>> "myParent" in cmds.ls()
         True
@@ -5769,7 +5803,7 @@ class DagModifier(_BaseModifier):
         >>> mod = DagModifier()
         >>> _ = mod.delete(child)
         >>> mod.doIt()
-        >>> parent.child() is None
+        >>> parent.child().name() == 'keepAlive'
         True
         >>> "myChild" in cmds.ls()
         False