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_dict = dict()
  28     rigs = []
  29     implementation_rigs = []
  30     MODULE_DIR = os.path.dirname(__file__)
  31     RIG_DIR_ABS = os.path.join(MODULE_DIR, utils.RIG_DIR)
  32     SEARCH_DIR_ABS = os.path.join(RIG_DIR_ABS, path)
  33     files = os.listdir(SEARCH_DIR_ABS)
  34     files.sort()
  35
  36     for f in files:
  37         is_dir = os.path.isdir(os.path.join(SEARCH_DIR_ABS, f))  # Whether the file is a directory
  38
  39         # Stop cases
  40         if f[0] in [".", "_"]:
  41             continue
  42         if f.count(".") >= 2 or (is_dir and "." in f):
  43             print("Warning: %r, filename contains a '.', skipping" % os.path.join(SEARCH_DIR_ABS, f))
  44             continue
  45
  46         if is_dir:
  47             # Check directories
  48             module_name = os.path.join(path, f).replace(os.sep, ".")
  49             rig = utils.get_rig_type(module_name)
  50             # Check if it's a rig itself
  51             if hasattr(rig, "Rig"):
  52                 rigs += [f]
  53             else:
  54                 # Check for sub-rigs
  55                 sub_dict = get_rig_list(os.path.join(path, f, ""))  # "" adds a final slash
  56                 rigs.extend(["%s.%s" % (f, l) for l in sub_dict['rig_list']])
  57                 implementation_rigs.extend(["%s.%s" % (f, l) for l in sub_dict['implementation_rigs']])
  58         elif f.endswith(".py"):
  59             # Check straight-up python files
  60             t = f[:-3]
  61             module_name = os.path.join(path, t).replace(os.sep, ".")
  62             rig = utils.get_rig_type(module_name)
  63             if hasattr(rig, "Rig"):
  64                 rigs += [t]
  65             if hasattr(rig, 'IMPLEMENTATION') and rig.IMPLEMENTATION:
  66                 implementation_rigs += [t]
  67     rigs.sort()
  68
  69     rigs_dict['rig_list'] = rigs
  70     rigs_dict['implementation_rigs'] = implementation_rigs
  71
  72     return rigs_dict
  73
  74
  75 def get_collection_list(rig_list):
  76     collection_list = []
  77     for r in rig_list:
  78         a = r.split(".")
  79         if len(a) >= 2 and a[0] not in collection_list:
  80             collection_list += [a[0]]
  81     return collection_list
  82
  83
  84 # Public variables
  85 rigs_dict = get_rig_list("")
  86 rig_list = rigs_dict['rig_list']
  87 implementation_rigs = rigs_dict['implementation_rigs']
  88 collection_list = get_collection_list(rig_list)
  89 col_enum_list = [("All", "All", ""), ("None", "None", "")] + [(c, c, "") for c in collection_list]