context_editor.py

import bpy
import warnings

bpy_context_key_types = dict([
    ("active_base",                     bpy.types.ObjectBase),
    ("active_bone",                     bpy.types.EditBone),
    ("active_gpencil_brush",            bpy.types.GPencilBrush),
    ("active_gpencil_frame",           [bpy.types.GPencilLayer]),
    ("active_gpencil_layer",           [bpy.types.GPencilLayer]),
    ("active_gpencil_palette",         [bpy.types.GPencilPalette]),
    ("active_gpencil_palettecolor",    [bpy.types.GPencilPaletteColor]),
    ("active_object",                   bpy.types.Object),
    ("active_operator",                 bpy.types.Operator),
    ("active_pose_bone",                bpy.types.PoseBone),
    ("area",                            bpy.types.Area),
    ("blend_data",                      bpy.types.BlendData),
    ("edit_object",                     bpy.types.Object),
    ("edit_text",                       bpy.types.Text),
    ("editable_bases",                 [bpy.types.ObjectBase]),
    ("editable_bones",                 [bpy.types.EditBone]),
    ("editable_gpencil_layers",        [bpy.types.GPencilLayer]),
    ("editable_gpdencil_strokes",      [bpy.types.GPencilStroke]),
    ("editable_objects",               [bpy.types.Object]),
    ("gpencil_data",                    bpy.types.GreasePencil), 
    ("gpencil_data_owner",              bpy.types.ID),
    ("image_paint_object",              bpy.types.Object),
    ("mode",                            str),
    ("object",                          bpy.types.Object),
    ("particle_edit_object",            bpy.types.Object),
    ("region",                          bpy.types.Region),
    ("region_data",                     bpy.types.RegionView3D),
    ("scene",                           bpy.types.Scene),
    ("screen",                          bpy.types.Screen),
    ("sculpt_object",                   bpy.types.Object),
    ("selectable_bases",               [bpy.types.ObjectBase]),
    ("selectable_objects",             [bpy.types.Object]),
    ("selected_bases",                 [bpy.types.ObjectBase]),
    ("selected_bones",                 [bpy.types.EditBone]),
    ("selected_editable_bases",        [bpy.types.ObjectBase]),
    ("selected_editable_bones",        [bpy.types.EditBone]),
    ("selected_editable_objects",      [bpy.types.Object]),
    ("selected_editable_sequences",    [bpy.types.Sequence]),
    ("selected_objects",               [bpy.types.Object]),
    ("selected_pose_bones",            [bpy.types.PoseBone]),
    ("selected_sequences",             [bpy.types.Sequence]),
    ("sequences",                      [bpy.types.Sequence]),
    ("space_data",                      bpy.types.Space),
    ("tool_settings",                   bpy.types.ToolSettings),
    ("user_preferences",                bpy.types.UserPreferences),
    ("vertex_paint_object",             bpy.types.Object),
    ("visible_bases",                  [bpy.types.ObjectBase]),
    ("visible_bones",                  [bpy.types.EditBone]),
    ("visible_gpencil_layers",         [bpy.types.GPencilLayer]),
    ("visible_objects",                [bpy.types.Object]),
    ("visible_pose_bones",             [bpy.types.PoseBone]),
    ("weight_paint_object",             bpy.types.Object),
    ("window",                          bpy.types.Window),
    ("window_manager",                  bpy.types.WindowManager)
])

inclusion_rels = [
    # rule : left >= right
    ## objects
    ("selectable_objects","selected_objects"),
    ("selected_objects","selected_editable_objects"),
    ("editable_objects", "selected_editable_objects"),
    ("selected_objects", "active_object"),
    # TODO: visible_objects?

    ## bases
    ("selectable_bases","selected_bases"),
    ("selected_bases","selected_editable_bases"),
    ("editable_bases", "selected_editable_bases"),
    ("selected_bases", "active_base"),
    # TODO: visible_bases

    ## edit bones
    ("selected_bones","selected_editable_bones"),
    ("editable_bones", "selected_editable_bones"),
    ("selected_bones", "active_bone"),
    # TODO: visible_bones

    ## pose bones
    ("selected_pose_bones", "active_pose_bone"),
    # TODO: visible_pose_bones

    # TODO: gpencils, sequences and others?
]
inclusion_rels_weak = [
    # rule : left >= right but only when len(left) > 0
    ## objects
    ("selected_objects", "active_object"),
    ("selected_bases", "active_base"),
    ("selected_bones", "active_bone"),
    ("selected_pose_bones", "active_pose_bone"),
    # TODO: gpencils, sequences and others?
]

limited_attr = dict([
    # key: func that returns True if it is in limitation
    ("edit_object", lambda item: item.type in ["MESH", "CURVE", "SURFACE", "FONT", "ARMATURE","LATTICE"])
])
class ContextEditor():
    """High-level bpy.context.copy() handler

    for example, setting a object for something will set object_base and other related keys
    
    Constructor : ContextEditor(raw_context, set_how)
        raw_context : dict | None
            = a dict that has the keys:values like bpy.context.copy()
              if None, it will automatically copy the current context
        how : str in {"AUTO", "MANUAL"}
            = a str to choose editing method
                if "AUTO":
                    All the related keys will be set automatically to avoid contradiction (such as "active_object not in selected_objects") 
                    This will ignore the actual context 
                    ex) all hidden objects are added to selectable_objects
                if "MANUAL":
                    None of the related keys are changed so contradiction may occur
    """
    original_keys = dict([
        ("area_type", str)])
    context_key_types = {**original_keys, **bpy_context_key_types}
    key_AUTO = ["selected_objects", "active_object", "object", "edit_object", "mode", "region", "screen", "window_manager", "window", "area_type"] 
    
    @property
    def raw_context(self):
        return self.consistant_dict(self._rc)
    @property
    def rc(self):
        return self.raw_context
    @property
    def dict_(self):
        rc = self.rc
        if self._how == "AUTO":
            return {k:rc[k] for k in self.key_AUTO}
        else:
            return rc 
    @property
    def how(self):
        return self._how

    def __init__(self, raw_context = None, how = "AUTO"):
        _raw_context = bpy.context.copy()
        if raw_context is not None:
            _raw_context.update(raw_context)
        for k in set(self.original_keys) - set(_raw_context):
            if k == "area_type":
                _raw_context[k] = _raw_context["area"].type
        self._rc = _raw_context.copy()
        self._how = how
        self.set_how(how)
    def set_how(self, how):
        if how == self._how:
            return
        self._rc = self.consistant_dict(self._rc)
        self._how = how
    @classmethod
    def consistant_dict(cls, raw_context):
        """amend contradictions of _raw_context"""
        rc = raw_context.copy()
        # set values based on inclusion_rels in sequences
        for big_k, small_k in inclusion_rels:
            small = rc[small_k]
            if small is None:
                continue
            big = set(rc[big_k])
            if len(big) == 0:
                if (big_k, small_k) in inclusion_rels_weak:
                    continue
            if type(small) == list:
                small = set(small)
            else:
                small = set([small])
            rc[big_k] += list(small - big)

        # set area based on area_type
        cls._make_ref()
        areas = [area for area in cls._areas 
                if area.type == rc["area_type"]]
        if areas:
            rc["area"] = areas[0]
        else:
            warnings.warn("Could not find area of {0}. area_type overwritten to {1}".format(rc["area_type"], rc["area"].type))
            rc["area_type"] = rc["area"].type
        rc["space_data"] = rc["area"].spaces[0]
        return rc
    def __setitem__(self, key, item):
        item = self._check_type(key, item)
        if self._how == "MANUAL":
            self._rc[key] = item
            return
        # only for AUTO 
        update_d = dict([]) # used for updating _raw_context
        self._make_ref()
        if key in ["selected_objects"]:
            selected_bases = [self._obj2base[obj.name] for obj in item]
            # objects
            update_d["selected_objects"] = item
            update_d["selected_bases"] = selected_bases
            update_d["selected_editable_objects"] = item
            update_d["selected_editable_bases"] = selected_bases
            if item: # do not update if item == []
                update_d["active_object"] = item[-1]
                update_d["object"] = item[-1]
                update_d["active_base"] = selected_bases[-1] 
        elif key in ["active_object", "object"]:
            # objects
            if item is not None:
                update_d["active_object"] = item
                update_d["object"] = item
                update_d["active_base"] = self._obj2base[item.name]
                update_d["selected_objects"] = self._rc["selected_objects"] + [item]
                update_d["selected_bases"] = self._rc["selected_bases"] + [update_d["active_base"]]
                update_d["selected_editable_objects"] = self._rc["selected_editable_objects"] + [item]
                update_d["selected_editable_bases"] = self._rc["selected_editable_bases"] + [update_d["active_base"]]
            else:
                update_d["active_object"] = update_d["active_base"] = None
        elif key in ["edit_object"]:
            # objects
            if item is not None:
                update_d["active_object"] = item
                update_d["active_base"] = self._obj2base[item.name]
            else:
                update_d["active_object"] = update_d["active_base"] = None
            mode = self._rc["mode"]
            if item is None: 
                if mode[:5] != "EDIT_":
                    update_d["mode"] = "OBJECT"
            elif mode in ["OBJECT"]:
                # TODO
                pass
        elif key in ["mode"]:
            update_d[key] = item
            e_obj = self._rc["edit_object"]
            if mode == "OBJECT":
                update_d["edit_object"] = None
            elif mode[:5] == "EDIT_" and e_obj is None:
                update_d["edit_object"] = self._rc["active_object"]
        else:
            update_d[key] = item
        self._rc.update(update_d)
        
    
    def _check_type(self,key,item):
        if key not in self.keys():
            raise IndexError("No such key as: {0}".format(key))
        type_ = self.context_key_types[key]
        if type(type_) != list:
            if item is None:
                return item
            if isinstance(item, type_):
                if key not in limited_attr.keys() or limited_attr["key"](item):
                    return item
            err_msg = "{0} only allows {1}".format(key,type_)
        else:
            # make type to list
            if isinstance(item, (list, tuple, set)):
                item = list(item)
            elif isinstance(item, dict):
                item = list(item.values())
            else: 
                item = [item]

            if len(item) == 0:
                return item # []
            type_ = type_[0]
            if all([isinstance(i, type_) for i in item]):
                if key not in limited_attr.keys() or limited_attr["key"](item):
                    return item
            err_msg = "{0} only allows sequences of {1}".format(key,type_)
        raise TypeError(err_msg)
    @classmethod
    def _make_ref(cls):
        """make references for data hard to access (like object_base)"""
        cls._bases = [base for scene in bpy.data.scenes 
                for base in scene.object_bases]
        cls._obj2base = dict(
                [(base.object.name, base) for base in cls._bases])
        cls._areas = [area for screen in bpy.data.screens
                            for area in screen.areas]

    def update(self, iterable):
        iterable = dict(iterable)
        for key, item in iterable.items():
            try:
                self.__setitem__(key, item)
            except (IndexError, ValueError) as e:
                warnings.warn(e)
        self._rc.update(iterable)

    def __getitem__(self, key):
        if type(self.dict_[key]) == list:
            return self.dict_[key][:]
        return self.dict_[key]

    # copy ordinary dict operations

    def __repr__(self):
        return repr(self.dict_)

    def __len__(self):
        return len(self.dict_)

    def __delitem__(self, key):
        if self._how == "AUTO":
            raise RuntimeError("You cannot delete any items in AUTO mode")
        else:
            del self.dict_[key]
            del self._rc[key]

    def clear(self):
        return self.dict_.clear()

    def copy(self):
        return self.dict_.copy()

    def has_key(self, k):
        return k in self.dict_

    def keys(self):
        return self.dict_.keys()

    def values(self):
        return self.dict_.values()

    def items(self):
        return self.dict_.items()

    def pop(self, *args):
        if self._how == "AUTO":
            raise RuntimeError("You cannot delete any items in AUTO mode")
        else:
            self._rc.pop(*args)
            return self.dict_.pop(*args)
    def __contains__(self, item):
        return item in self.dict_

    def __iter__(self):
        return iter(self.dict_)

def main():
    c = ContextEditor()
    c["selected_objects"] = []
    c["object"] = bpy.data.objects["Cube.001"]
    bpy.ops.object.origin_set(c.rc,type='ORIGIN_GEOMETRY')
if __name__ == "__main__":
    main()