1    from __future__ import print_function, division
2    import sys
3    
4    if sys.version_info.major > 2:
5        basestring = str
6    else:
7        range = xrange
8    
9    try:
10       import poser
11   except ImportError:
12       import POSER_FAKE as poser
13   
14   SCENE = poser.Scene()
15   ROW_HEIGHT = 150
16   
17   
18   def mat_layernames(material):
19       """ 
20       Generator. 
21       Yields all layernames. 
22       """
23       assert isinstance(material, poser.MaterialType)
24       for layer in material.Layers():
25           yield layer.ExtName()
26   
27   
28   def get_stree_from_layer(material, layer_name):
29       """ 
30       Returns shadertree from a layer given its name. 
31       Returns None if there is no layer with this name. 
32    
33       :type material: poser.MaterialType 
34       :type layer_name: basestring 
35       :rtype: poser.ShaderTree 
36       """
37       assert isinstance(layer_name, basestring)
38       for layer in material.Layers():
39           if layer.ExtName() == layer_name:
40               return layer.ShaderTree()
41       return None
42   
43   
44   def force_stree_from_layer(material, layer_name):
45       """ 
46       Get shadertree from layer by name (see above). 
47       But creates the layer if it does not exists. 
48    
49       :type layer_name: basestring 
50       :type material: poser.MaterialType 
51       :rtype: poser.ShaderTree 
52       """
53       if layer_name in mat_layernames(material):
54           return get_stree_from_layer(material, layer_name)
55   
56       material.CreateLayer(layer_name)
57       return material.LayerByExtName(layer_name).ShaderTree()
58   
59   
60   class ShaderNode(object):
61       """ 
62       Extends a Poser ShaderNode (poser.ShaderNodeType). 
63       Most important: One may use it like a dictionary with fixed keys: 
64           # Say we want to create a Poser Bright/Contrast-Node 
65           # (assumed we have a ShaderTree; see below): 
66           stree = ShaderTree(material.ShaderTree()) 
67           node = stree.newnode(poser.kNodeTypeCodeCyclesBRIGHTCONTRAST, 
68                                Color=rgb(200, 250, 16), 
69                                Bright=0, 
70                                Contrast=0) 
71       After this the new node is created and shown in the Materialroom. 
72       All inputs are set as given above. 
73    
74       The node can be manipulated anytime. To set a single attribute: 
75           node["Bright"] = 0.2 
76       To set multiple attributes at once: 
77           node(Color=grayscale(0.98), Bright=1.2, Contrast=0.1) 
78       For the first and last method one may need a possibility to deal 
79       with blank characters for some attributes. If a blank is needed 
80       in an inputname, use double underscores: Blend__Type=... for "Blend Type". 
81    
82       The order you declare the attribute names is not important. No 
83       problem to declare the last first and the first last. 
84    
85       To get a value from an attribute: 
86           v = node["Bright"]  # returns 0.2 
87    
88       A check if an attributenames exist is also possible: 
89           if "Bright" in node: 
90               ... 
91    
92       Lots of attributes may not have only one but two or three values. 
93       For this, there is a rule of thumbs: 
94           One numeric value: Numeric float. 
95           One string value: String, mostly a filename. 
96           One array with 3 floats: Color value (R, G, B). 
97           One NodeType: A node to connect. 
98    
99       To set different values you may do this: 
100          texture_node = stree.newnode(poser.kNodeTypeCodeImageMap, ...) 
101          node["Color"] = (r, g, b), texture_node 
102           - or - 
103          node["Color"] = texture_node, (r, g, b) 
104   
105      Means: To assign different types of inputs, the ShaderNode class checks 
106      for different valuetypes to sets the right one. So it can automatically 
107      find out how a value must be dealt with. 
108      A check will be made if a tuple/list is given. Has a tuple/list 3 elements 
109      is it a color value (r, g, b). Has it only two elements, a pair of 
110      value and another node is (should be) meant. 
111   
112      If another node is assigned to an attribute, this other node is connected to 
113      the named input. 
114   
115      A Brightness/Contrast node again: 
116          stree.newnode(poser.kNodeTypeCodeCyclesBRIGHTCONTRAST, 
117                        Color=(stree.newnode(poser.kNodeTypeCodeImageMap, 
118                                             Image="MyImage.png", 
119                                             ), 
120                               rgb(250, 250, 250)), 
121                        Bright=0, 
122                        Contrast=0) 
123      Creates the Bright/Contrast node like above and attaches a new 
124      created imagemap and a color entry to the Color attribute in one go. 
125      Note that the new created node is not assigned to a variable. That's 
126      ok, if you don't need a reference for later use. 
127   
128      All nodes have some special attributes: 
129          Name: String, Exposed nodename. 
130          pos:  tuple(int x, int y), screen position. 
131          collapsed: bool. 
132   
133      See Example for more. 
134      """
135  
136      def __init__(self, poser_node, **kwargs):
137          if isinstance(poser_node, ShaderNode):
138              # Handle nodes already wrapped in ShaderNode.
139              poser_node = poser_node.node
140  
141          assert isinstance(poser_node, poser.ShaderNodeType)
142  
143          # If a parent is given, it must be a shaderrtree.
144          # Attention: Some functions may need this parent, others not.
145          self.__parent = kwargs.pop("parent", kwargs.pop("Parent", None))
146          assert isinstance(self.__parent, ShaderTree) or self.__parent is None
147  
148          self.node = poser_node
149          # Get the names from this nodes attributes.
150          self.__input_names = [_inp.Name() for _inp in poser_node.Inputs()]
151          self.__input_inames = [_inp.InternalName() for _inp in poser_node.Inputs()]
152  
153          # Assign values given in kwargs (if any).
154          for k, v in kwargs.items():
155              self[k.replace("__", " ")] = v
156  
157      def __getattr__(self, item):
158          # Any attribute this class does not know goes to the
159          # original Poser node.
160          if hasattr(self.node, item):
161              return getattr(self.node, item)
162          return super(ShaderNode, self).__getattribute__(item)
163  
164      def __contains__(self, item):
165          return item in self.__input_names or item in self.__input_inames
166  
167      def __getitem__(self, item):
168          # Return a poser.ShaderNodeInputType
169          return self.Input(item)
170  
171      def __setitem__(self, key, value):
172          # Handle the possible value types.
173          def do_set(_v=value):
174              if isinstance(_v, (int, float)):
175                  inp.SetFloat(_v)
176              elif isinstance(_v, basestring):
177                  inp.SetString(_v)
178              elif isinstance(_v, (list, tuple)) and len(_v) == 3:
179                  inp.SetColor(*_v)
180              elif isinstance(_v, (ShaderNode, poser.ShaderNodeType)):
181                  _v.ConnectToInput(inp)
182  
183          if key.startswith("pos"):
184              if key == "pos":
185                  self.node.SetLocation(*value)
186              elif key == "posX":
187                  p = self.node.Location()
188                  self.node.SetLocation(value, p[1])
189              elif key == "posY":
190                  p = self.node.Location()
191                  self.node.SetLocation(p[0], value)
192          elif key == "collapse":
193              self.parent.collapse = value
194          elif key == "Name":
195              self.node.SetName(value)
196          else:
197              inp = self.Input(key)
198              if inp:
199                  if isinstance(value, (list, tuple)) and len(value) == 2:
200                      do_set(value[0])
201                      do_set(value[1])
202                  else:
203                      do_set(value)
204  
205      def __call__(self, *args, **kwargs):
206          for k, v in kwargs.items():
207              self[k.replace("__", " ")] = v
208          return self
209  
210      @property
211      def parent(self):
212          """:rtype: ShaderTree"""
213          return self.__parent
214  
215      def Input(self, arg):
216          """ 
217          Overwrite Input from original node with an extended version. 
218          If an integer value is given, it works like the original. 
219          But instead of a number one may also provide an attributename. 
220          """
221          if isinstance(arg, basestring):
222              idx = -1
223              if arg in self.__input_names:
224                  idx = self.__input_names.index(arg)
225              elif arg in self.__input_inames:
226                  idx = self.__input_inames.index(arg)
227              if idx != -1:
228                  return self.node.Input(idx)
229  
230              raise AttributeError("Node '%s' has no '%s'." % (self.node.Name(), arg))
231  
232          return self.node.Input(arg)
233  
234      def input_dict(self):
235          """ 
236          Returns all attributes from the node as a dictionary. 
237          :rtype: dict() 
238          """
239          d = dict(Name=self.node.Name(), pos=list(self.node.Location()))
240          for iname in self.__input_inames:
241              d[iname] = self[iname]
242          return d
243  
244      def copy(self, from_node):
245          """ 
246          Copy all attributes from "from_node" to this node. 
247          Raises TypeError if both nodes are not the same type. 
248          """
249          assert isinstance(from_node, (ShaderNode, poser.ShaderNodeType))
250          if self.Type() != from_node.Type():
251              raise TypeError("Nodes must be of same type.")
252  
253          self.node.SetName(from_node.Name())
254          # self.node.SetLocation(from_node.Location())
255          self.node.SetPreviewCollapsed(from_node.PreviewCollapsed())
256          self.node.SetInputsCollapsed(from_node.InputsCollapsed())
257          self.__parent = from_node.__parent
258          for iname in self.__input_inames:
259              self[iname] = from_node[iname]
260  
261  
262  class ShaderTree(object):
263      """ 
264      Modified version of a Poser Shadertree. 
265   
266      If a node is returned it is always wrapped in the above class "ShaderNode". 
267      """
268  
269      def __init__(self, poser_tree, Parent=None):
270          if isinstance(poser_tree, ShaderTree):
271              poser_tree = poser_tree.tree
272          assert isinstance(poser_tree, poser.ShaderTreeType)
273          self.tree = poser_tree
274          self._collapse_default = False
275          self.__parent = Parent  # type: poser.MaterialType
276  
277      def __getattr__(self, item):
278          if hasattr(self.tree, item):
279              res = getattr(self.tree, item)
280          else:
281              res = super(ShaderTree, self).__getattribute__(item)
282          if isinstance(res, poser.ShaderNodeType):
283              return ShaderNode(res)
284          return res
285  
286      @property
287      def parent(self):
288          return self.__parent
289  
290      @property
291      def collapse(self):
292          return self._collapse_default
293  
294      @collapse.setter
295      def collapse(self, v):
296          self._collapse_default = bool(v)
297  
298      def create_node(self, node_type, **kwargs):
299          node = self.tree.CreateNode(node_type)
300          if node is None:
301              return
302  
303          collapse = kwargs.pop("collapse", self._collapse_default)
304          if collapse:
305              node.SetInputsCollapsed(1)
306              node.SetPreviewCollapsed(0)
307          else:
308              node.SetInputsCollapsed(0)
309              node.SetPreviewCollapsed(1)
310  
311          node.SetName(kwargs.pop("Name", node.Name()))
312          x, y = kwargs.pop("pos", node.Location())
313          node.SetLocation(kwargs.pop("posX", x), kwargs.pop("posY", y))
314          return ShaderNode(node, Parent=self, **kwargs)
315  
316      newnode = create_node
317  
318      def root(self, renderertype):
319          n = self.tree.RendererRootNode(renderertype)
320          return ShaderNode(n, Parent=self) if n else None
321  
322      @property
323      def cycles_root(self):
324          """:rtype: ShaderNode"""
325          root = self.root(poser.kRenderEngineCodeSUPERFLY)
326          if root is None or root.Type() != "CyclesSurface":
327              root = ShaderNode(self.tree.CreateNode("CyclesSurface"), Parent=self)
328              self.tree.SetRendererRootNode(poser.kRenderEngineCodeSUPERFLY, root.node)
329          return root
330  
331      @property
332      def pysical_root(self):
333          """:rtype: ShaderNode"""
334          root = self.root(poser.kRenderEngineCodeSUPERFLY)
335          if root is None or root.Type() != "PhysicalSurface":
336              root = self.tree.CreateNode("PhysicalSurface")
337              self.tree.SetRendererRootNode(poser.kRenderEngineCodeSUPERFLY, root.node)
338          return root
339  
340      @property
341      def firefly_root(self):
342          """:rtype: ShaderNode"""
343          root = self.root(poser.kRenderEngineCodeFIREFLY)
344          if root is None or root.Type() != "poser":
345              root = self.tree.CreateNode("poser")
346              self.tree.SetRendererRootNode(poser.kRenderEngineCodeSUPERFLY, root.node)
347          return root
348  
349      def remove_surface(self, code):
350          for node in self.tree.Nodes():
351              if node.Type() == code:
352                  node.Delete()
353  
354      def clear_tree(self):
355          for node in self.tree.Nodes():  # type: poser.ShaderNodeType
356              node.Delete()
357  
358      def find_node(self, displayname):
359          """ 
360          :param displayname: 
361          :rtype: poser.ShaderNodeType 
362          """
363          for node in self.tree.Nodes():
364              if node.Name() == displayname or node.InternalName() == displayname:
365                  return ShaderNode(node, Parent=self)
366          return None
367  
368      def yield_node_types(self, nodetypes):
369          if not isinstance(nodetypes, (list, tuple)):
370              nodetypes = [nodetypes, ]
371  
372          for node in self.tree.Nodes():
373              if node.Type() in nodetypes:
374                  yield ShaderNode(node, Parent=self)
375  
376      def yield_images(self):
377          for img_node in self.yield_node_types((poser.kNodeTypeCodeIMAGEMAP,
378                                                 poser.kNodeTypeCodeCyclesIMAGETEXTURE)):
379              yield img_node
380  
381      def find_image(self, displayname=None, filename=None):
382          if displayname is not None:
383              return self.find_node(displayname)
384  
385          if filename is not None:
386              for img in self.yield_images():
387                  if "Image" in img and img["Image"] == filename:
388                      return img
389                  if "Image_Source" in img and img["Image_Source"] == filename:
390                      return img
391  
392      def get_images(self):
393          return list(self.yield_images())
394