rigify/rig_lists.py

   1 #====================== BEGIN GPL LICENSE BLOCK ======================
   2 #
   3 #  This program is free software; you can redistribute it and/or
   4 #  modify it under the terms of the GNU General Public License
   5 #  as published by the Free Software Foundation; either version 2
   6 #  of the License, or (at your option) any later version.
   7 #
   8 #  This program is distributed in the hope that it will be useful,
   9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
  10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  11 #  GNU General Public License for more details.
  12 #
  13 #  You should have received a copy of the GNU General Public License
  14 #  along with this program; if not, write to the Free Software Foundation,
  15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
  16 #
  17 #======================= END GPL LICENSE BLOCK ========================
  18
  19 import os
  20
  21 from . import utils
  22
  23
  24 def get_rig_list(path):
  25     """ Recursively searches for rig types, and returns a list.
  26     """
  27     rigs = []
  28     MODULE_DIR = os.path.dirname(__file__)
  29     RIG_DIR_ABS = os.path.join(MODULE_DIR, utils.RIG_DIR)
  30     SEARCH_DIR_ABS = os.path.join(RIG_DIR_ABS, path)
  31     files = os.listdir(SEARCH_DIR_ABS)
  32     files.sort()
  33
  34     for f in files:
  35         is_dir = os.path.isdir(os.path.join(SEARCH_DIR_ABS, f))  # Whether the file is a directory
  36
  37         # Stop cases
  38         if f[0] in [".", "_"]:
  39             continue
  40         if f.count(".") >= 2 or (is_dir and "." in f):
  41             print("Warning: %r, filename contains a '.', skipping" % os.path.join(SEARCH_DIR_ABS, f))
  42             continue
  43
  44         if is_dir:
  45             # Check directories
  46             module_name = os.path.join(path, f).replace(os.sep, ".")
  47             rig = utils.get_rig_type(module_name)
  48             # Check if it's a rig itself
  49             if hasattr(rig, "Rig"):
  50                 rigs += [f]
  51             else:
  52                 # Check for sub-rigs
  53                 ls = get_rig_list(os.path.join(path, f, ""))  # "" adds a final slash
  54                 rigs.extend(["%s.%s" % (f, l) for l in ls])
  55         elif f.endswith(".py"):
  56             # Check straight-up python files
  57             t = f[:-3]
  58             module_name = os.path.join(path, t).replace(os.sep, ".")
  59             rig = utils.get_rig_type(module_name)
  60             if hasattr(rig, "Rig"):
  61                 rigs += [t]
  62     rigs.sort()
  63     return rigs
  64
  65
  66 def get_collection_list(rig_list):
  67     collection_list = []
  68     for r in rig_list:
  69         a = r.split(".")
  70         if len(a) >= 2 and a[0] not in collection_list:
  71             collection_list += [a[0]]
  72     return collection_list
  73
  74
  75 # Public variables
  76 rig_list = get_rig_list("")
  77 collection_list = get_collection_list(rig_list)
  78 col_enum_list = [("All", "All", ""), ("None", "None", "")] + [(c, c, "") for c in collection_list]