space_view3d_math_vis/utils.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 # <pep8 compliant>
  20
  21 import bpy
  22
  23
  24 def console_namespace():
  25     import console_python
  26     get_consoles = console_python.get_console
  27     consoles = getattr(get_consoles, "consoles", None)
  28     if consoles:
  29         for console, stdout, stderr in get_consoles.consoles.values():
  30             return console.locals
  31     return {}
  32
  33
  34 def is_display_list(listvar):
  35     from mathutils import Vector
  36
  37     for var in listvar:
  38         if type(var) is not Vector:
  39             return False
  40     return True
  41
  42
  43 class VarStates:
  44
  45     @staticmethod
  46     def store_states():
  47         # Store the display states, called upon unregister the Add-on
  48         # This is useful when you press F8 to reload the Addons.
  49         # Then this function preserves the display states of the
  50         # console variables.
  51         state_props = bpy.context.window_manager.MathVisStatePropList
  52         variables = get_math_data()
  53         for key, ktype in variables.items():
  54             if key and key not in state_props:
  55                 prop = state_props.add()
  56                 prop.name = key
  57                 prop.ktype = ktype.__name__
  58                 prop.state = [True, False]
  59
  60     @staticmethod
  61     def get_index(key):
  62         index = bpy.context.window_manager.MathVisStatePropList.find(key)
  63         return index
  64
  65     @staticmethod
  66     def delete(key):
  67         state_props = bpy.context.window_manager.MathVisStatePropList
  68         index = state_props.find(key)
  69         if index != -1:
  70             state_props.remove(index)
  71
  72     @staticmethod
  73     def toggle_display_state(key):
  74         state_props = bpy.context.window_manager.MathVisStatePropList
  75         if key in state_props:
  76             state_props[key].state[0] = not state_props[key].state[0]
  77         else:
  78             print("Odd: Can not find key %s in MathVisStateProps" % (key))
  79
  80     @staticmethod
  81     def toggle_lock_state(key):
  82         state_props = bpy.context.window_manager.MathVisStatePropList
  83         if key in state_props:
  84             state_props[key].state[1] = not state_props[key].state[1]
  85         else:
  86             print("Odd: Can not find key %s in MathVisStateProps" % (key))
  87
  88
  89 def get_math_data():
  90     from mathutils import Matrix, Vector, Quaternion, Euler
  91
  92     locals = console_namespace()
  93     if not locals:
  94         return {}
  95
  96     variables = {}
  97     for key, var in locals.items():
  98         if len(key) == 0 or key[0] == "_":
  99             continue
 100
 101         type_var = type(var)
 102
 103         # Rules out sets/dicts.
 104         # It's also possible the length check below is slow
 105         # for data with underlying linked-list structure.
 106         if not hasattr(type_var, "__getitem__"):
 107             continue
 108
 109         # Don't do a truth test on the data because this causes an error with some
 110         # array types, see T66107.
 111         len_fn = getattr(type_var, "__len__", None)
 112         if len_fn is None:
 113             continue
 114         if len_fn(var) == 0:
 115             continue
 116
 117         if type_var in {Matrix, Vector, Quaternion, Euler} or \
 118            type_var in {tuple, list} and is_display_list(var):
 119
 120             variables[key] = type_var
 121
 122     return variables
 123
 124
 125 def cleanup_math_data():
 126
 127     locals = console_namespace()
 128     if not locals:
 129         return
 130
 131     variables = get_math_data()
 132
 133     for key in variables.keys():
 134         index = VarStates.get_index(key)
 135         if index == -1:
 136             continue
 137
 138         state_prop = bpy.context.window_manager.MathVisStatePropList.get(key)
 139         if state_prop.state[1]:
 140             continue
 141
 142         del locals[key]
 143         bpy.context.window_manager.MathVisStatePropList.remove(index)
 144
 145
 146 def console_math_data():
 147     from mathutils import Matrix, Vector, Quaternion, Euler
 148
 149     data_matrix = {}
 150     data_quat = {}
 151     data_euler = {}
 152     data_vector = {}
 153     data_vector_array = {}
 154
 155     for key, var in console_namespace().items():
 156         if key[0] == "_":
 157             continue
 158
 159         state_prop = bpy.context.window_manager.MathVisStatePropList.get(key)
 160         if state_prop:
 161             disp, lock = state_prop.state
 162             if not disp:
 163                 continue
 164
 165         var_type = type(var)
 166
 167         if var_type is Matrix:
 168             if len(var.col) != 4 or len(var.row) != 4:
 169                 if len(var.col) == len(var.row):
 170                     var = var.to_4x4()
 171                 else:  # todo, support 4x3 matrix
 172                     continue
 173             data_matrix[key] = var
 174         elif var_type is Vector:
 175             if len(var) < 3:
 176                 var = var.to_3d()
 177             data_vector[key] = var
 178         elif var_type is Quaternion:
 179             data_quat[key] = var
 180         elif var_type is Euler:
 181             data_euler[key] = var
 182         elif var_type in {list, tuple} and is_display_list(var):
 183             data_vector_array[key] = var
 184
 185     return data_matrix, data_quat, data_euler, data_vector, data_vector_array