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 import traceback
  21 import importlib
  22
  23 from .utils.rig import RIG_DIR
  24
  25 from . import feature_set_list
  26
  27
  28 def get_rigs(base_dir, base_path, *, path=[], feature_set=feature_set_list.DEFAULT_NAME):
  29     """ Recursively searches for rig types, and returns a list.
  30
  31     :param base_path: base dir where rigs are stored
  32     :type path:str
  33     :param path:      rig path inside the base dir
  34     :type path:str
  35     """
  36
  37     rigs = {}
  38     impl_rigs = {}
  39
  40     dir_path = os.path.join(base_dir, *path)
  41
  42     try:
  43         files = os.listdir(dir_path)
  44     except FileNotFoundError:
  45         files = []
  46
  47     files.sort()
  48
  49     for f in files:
  50         is_dir = os.path.isdir(os.path.join(dir_path, f))  # Whether the file is a directory
  51
  52         # Stop cases
  53         if f[0] in [".", "_"]:
  54             continue
  55         if f.count(".") >= 2 or (is_dir and "." in f):
  56             print("Warning: %r, filename contains a '.', skipping" % os.path.join(*base_path, *path, f))
  57             continue
  58
  59         if is_dir:
  60             # Check for sub-rigs
  61             sub_rigs, sub_impls = get_rigs(base_dir, base_path, path=[*path, f], feature_set=feature_set)
  62             rigs.update(sub_rigs)
  63             impl_rigs.update(sub_impls)
  64         elif f.endswith(".py"):
  65             # Check straight-up python files
  66             subpath = [*path, f[:-3]]
  67             key = '.'.join(subpath)
  68             # Don't reload rig modules - it breaks isinstance
  69             rig_module = importlib.import_module('.'.join(base_path + subpath))
  70             if hasattr(rig_module, "Rig"):
  71                 rigs[key] = {"module": rig_module,
  72                              "feature_set": feature_set}
  73             if hasattr(rig_module, 'IMPLEMENTATION') and rig_module.IMPLEMENTATION:
  74                 impl_rigs[key] = rig_module
  75
  76     return rigs, impl_rigs
  77
  78
  79 # Public variables
  80 rigs = {}
  81 implementation_rigs = {}
  82
  83 def get_internal_rigs():
  84     global rigs, implementation_rigs
  85
  86     BASE_RIGIFY_DIR = os.path.dirname(__file__)
  87     BASE_RIGIFY_PATH = __name__.split('.')[:-1]
  88
  89     rigs, implementation_rigs = get_rigs(os.path.join(BASE_RIGIFY_DIR, RIG_DIR), [*BASE_RIGIFY_PATH, RIG_DIR])
  90
  91 def get_external_rigs(set_list):
  92     # Clear and fill rigify rigs and implementation rigs public variables
  93     for rig in list(rigs.keys()):
  94         if rigs[rig]["feature_set"] != feature_set_list.DEFAULT_NAME:
  95             rigs.pop(rig)
  96             if rig in implementation_rigs:
  97                 implementation_rigs.pop(rig)
  98
  99     # Get external rigs
 100     for feature_set in set_list:
 101         try:
 102             base_dir, base_path = feature_set_list.get_dir_path(feature_set, RIG_DIR)
 103
 104             external_rigs, external_impl_rigs = get_rigs(base_dir, base_path, feature_set=feature_set)
 105         except Exception:
 106             print("Rigify Error: Could not load feature set '%s' rigs: exception occurred.\n" % (feature_set))
 107             traceback.print_exc()
 108             print("")
 109             continue
 110
 111         rigs.update(external_rigs)
 112         implementation_rigs.update(external_impl_rigs)