add_camera_rigs: refactor and cleanup

- Fix widgets’ names: they were hardcoded and didn’t follow the
  preferences, leading to crashes.
- The UI was put back into the Item category, instead of Create,
  because it is not related to object creation.
- Fix some strange topology in two widget shapes.
- UI and operators use a new poll method, so that they work when
  either the rig or the camera is selected.
- The composition guides UI was converted to a panel, so that they may
  be drag-selected.
- Marker binding and DOF object operators were converted to the
  `bpy.data` API, making them simpler.
- Bones were moved around so that they are more similar between rigs.
  - They were scaled down to be 1 unit long, a simpler length — for
    instance, widgets are the same size as modeled. Widgets were
    scaled up to compensate.
  - The camera and aim bones were placed at 1.7 unit high, to be
    approximately at a standing human’s eyes’ height if the scene is
    in meters.
- Much of the rig generation was refactored to deduplicate code
  between the two rig types.
- Automatic renaming to `.000` was removed, since Blender already
  handles duplicate names.
- Widget prefix and collection were renamed to `WGT-` and `Widgets`
  respectively. This is to be closer to Rigify, hopefully unifying
  them.
- The GPL license header was added to every file.
- Some cleanup was done to better respect Python’s PEP 8.

Reviewed By: Wayne Dixon

Differential Revision: https://developer.blender.org/D6543

1 files changed, 200 insertions, 445 deletions

diff --git a/add_camera_rigs/build_rigs.py b/add_camera_rigs/build_rigs.py
index c7303198..1a37bc6e 100644
--- a/add_camera_rigs/build_rigs.py
+++ b/add_camera_rigs/build_rigs.py
@@ -1,6 +1,25 @@
+# ##### BEGIN GPL LICENSE BLOCK #####
+#
+#  This program is free software; you can redistribute it and/or
+#  modify it under the terms of the GNU General Public License
+#  as published by the Free Software Foundation; either version 2
+#  of the License, or (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program; if not, write to the Free Software Foundation,
+#  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+#
+# ##### END GPL LICENSE BLOCK #####
+
 import bpy
+from bpy_extras import object_utils
 from bpy.types import Operator
-from math import radians
+from math import radians, pi
 from rna_prop_ui import rna_idprop_ui_prop_get
 from .create_widgets import (create_root_widget,
                              create_widget,
@@ -9,518 +28,254 @@ from .create_widgets import (create_root_widget,
                              )
 
 
-def build_dolly_rig(context):
-    """Operator to build the dolly rig"""
-    # Set the bone layers
-    boneLayer = (False, True, False, False, False, False, False, False,
-                 False, False, False, False, False, False, False, False,
-                 False, False, False, False, False, False, False, False,
-                 False, False, False, False, False, False, False, False)
-
-    # Add the new armature object:
-    bpy.ops.object.armature_add()
-    rig = context.active_object
-
-    # it will try to name the rig "Dolly_rig" but if that name exists it will
-    # add 000 to the name
-    if "Dolly_rig" not in context.scene.objects:
-        rig.name = "Dolly_rig"
-    else:
-        rig.name = "Dolly_rig.000"
-    rig["rig_id"] = "Dolly_rig"
-
-    bpy.ops.object.mode_set(mode='EDIT')
-
-    # Remove default bone:
+def create_dolly_bones(rig, bone_layers):
     bones = rig.data.edit_bones
-    bones.remove(bones[0])
 
-    # Add new bones:
+    # Add new bones
     root = bones.new("Root")
-    root.tail = (0.0, 3.0, 0.0)
+    root.tail = (0.0, 1.0, 0.0)
+    root.show_wire = True
 
-    bpy.ops.object.mode_set(mode='EDIT')
-    ctrlAimChild = bones.new("aim_MCH")
-    ctrlAimChild.head = (0.0, 5.0, 3.0)
-    ctrlAimChild.tail = (0.0, 7.0, 3.0)
-    ctrlAimChild.layers = boneLayer
+    ctrl_aim_child = bones.new("Aim_shape_rotation-MCH")
+    ctrl_aim_child.head = (0.0, 10.0, 1.7)
+    ctrl_aim_child.tail = (0.0, 11.0, 1.7)
+    ctrl_aim_child.layers = bone_layers
 
-    ctrlAim = bones.new("Aim")
-    ctrlAim.head = (0.0, 5.0, 3.0)
-    ctrlAim.tail = (0.0, 7.0, 3.0)
+    ctrl_aim = bones.new("Aim")
+    ctrl_aim.head = (0.0, 10.0, 1.7)
+    ctrl_aim.tail = (0.0, 11.0, 1.7)
+    ctrl_aim.show_wire = True
 
     ctrl = bones.new("Camera")
-    ctrl.head = (0.0, 0.0, 3.0)
-    ctrl.tail = (0.0, 2.0, 3.0)
+    ctrl.head = (0.0, 0.0, 1.7)
+    ctrl.tail = (0.0, 1.0, 1.7)
+    ctrl.show_wire = True
 
-    # Setup hierarchy:
+    # Setup hierarchy
     ctrl.parent = root
-    ctrlAim.parent = root
-    ctrlAimChild.parent = ctrlAim
-
-    # jump into pose mode and change bones to euler
-    bpy.ops.object.mode_set(mode='POSE')
-    for x in bpy.context.object.pose.bones:
-        x.rotation_mode = 'XYZ'
-
-    # jump into pose mode and add the custom bone shapes
-    bpy.ops.object.mode_set(mode='POSE')
-    bpy.context.object.pose.bones["Root"].custom_shape = bpy.data.objects[
-        "WDGT_Camera_root"]  # add the widget as custom shape
-    # set the wireframe checkbox to true
-    bpy.context.object.data.bones["Root"].show_wire = True
-    bpy.context.object.pose.bones[
-        "Aim"].custom_shape = bpy.data.objects["WDGT_Aim"]
-    bpy.context.object.data.bones["Aim"].show_wire = True
-    bpy.context.object.pose.bones["Aim"].custom_shape_transform = bpy.data.objects[
-        rig.name].pose.bones["aim_MCH"]  # sets the "At" field to the child
-    bpy.context.object.pose.bones[
-        "Camera"].custom_shape = bpy.data.objects["WDGT_Camera"]
-    bpy.context.object.data.bones["Camera"].show_wire = True
-
-    # jump into object mode
-    bpy.ops.object.mode_set(mode='OBJECT')
-
-    # Add constraints to bones:
-    con = rig.pose.bones['aim_MCH'].constraints.new('COPY_ROTATION')
-    con.target = rig
-    con.subtarget = "Camera"
-
-    con = rig.pose.bones['Camera'].constraints.new('TRACK_TO')
-    con.target = rig
-    con.subtarget = "Aim"
-    con.use_target_z = True
-
-    # Add custom Bone property to Camera bone
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "lock", create=True)
-    ob["lock"] = 1.0
-    prop["soft_min"] = prop["min"] = 0.0
-    prop["soft_max"] = prop["max"] = 1.0
-
-    # Add Driver to Lock/Unlock Camera from Aim Target
-    rig = bpy.context.view_layer.objects.active
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-
-    constraint = pose_bone.constraints["Track To"]
-    inf_driver = constraint.driver_add('influence')
-    inf_driver.driver.type = 'SCRIPTED'
-    var = inf_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["lock"]'
-    inf_driver.driver.expression = 'var'
-
-    # Add custom property for the lens / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "focal_length", create=True)
-    ob["focal_length"] = 50.0
-    prop["soft_min"] = prop["min"] = 1.0
-    prop["default"] = 50.0
-    prop["soft_max"] = prop["max"] = 5000.0
-
-    # Add custom property for the focus distance / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "focus_distance", create=True)
-    ob["focus_distance"] = 10.00
-    prop["soft_min"] = prop["min"] = 0.0
-    prop["soft_max"] = prop["max"] = 1000.0
-
-    # Add custom property for the f-stop / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "f-stop", create=True)
-    ob["f-stop"] = 2.8
-    prop["soft_min"] = prop["min"] = 0.1
-    prop["soft_max"] = prop["max"] = 128.00
-
-    # Add the camera object:
-    bpy.ops.object.mode_set(mode='OBJECT')
-
-    bpy.ops.object.camera_add()
-    cam = bpy.context.active_object
-
-    # Name the Camera Object
-    if 'Dolly_camera' not in context.scene.objects:
-        cam.name = "Dolly_camera"
-    else:
-        cam.name = "Dolly_camera.000"
-
-    # this will name the camera data
-    cam.data.name = cam.name
-
-    cam_data_name = bpy.context.object.data.name
-    bpy.data.cameras[cam_data_name].display_size = 1.0
-    cam.rotation_euler = [radians(90), 0, 0]  # rotate the camera 90 degrees in x
-
-    cam.location = (0.0, -2.0, 0.0)  # move the camera to the correct postion
-    cam.parent = rig
-    cam.parent_type = "BONE"
-    cam.parent_bone = "Camera"
-
-    # Add Driver to link the camera lens to the custom property on the armature
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-    lens_driver = cam.data.driver_add("lens")
-    lens_driver.driver.type = 'SCRIPTED'
-    var = lens_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["focal_length"]'
-    lens_driver.driver.expression = 'var'
-
-    # Add Driver to link the camera focus distance to the custom property on the armature
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-    lens_driver = cam.data.driver_add("dof.focus_distance")
-    lens_driver.driver.type = 'SCRIPTED'
-    var = lens_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["focus_distance"]'
-    lens_driver.driver.expression = 'var'
-
-    # Add Driver to link the camera f-stop to the custom property on the armature
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-    lens_driver = cam.data.driver_add("dof.aperture_fstop")
-    lens_driver.driver.type = 'SCRIPTED'
-    var = lens_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["f-stop"]'
-    lens_driver.driver.expression = 'var'
-
-    # lock the location/rotation/scale of the camera
-    cam.lock_location = [True, True, True]
-    cam.lock_rotation = [True, True, True]
-    cam.lock_scale = [True, True, True]
-
-    # Set new camera as active camera
-    bpy.context.scene.camera = cam
-
-    # make sure the camera is selectable by default (this can be locked in the UI)
-    bpy.context.object.hide_select = False
-
-    # make the rig the active object before finishing
-    bpy.context.view_layer.objects.active = rig
-    bpy.data.objects[cam.name].select_set(False)
-    bpy.data.objects[rig.name].select_set(True)
-
-    return rig
-
+    ctrl_aim.parent = root
+    ctrl_aim_child.parent = ctrl_aim
 
-class ADD_CAMERA_RIGS_OT_build_dolly_rig(Operator):
-    bl_idname = "add_camera_rigs.build_dolly_rig"
-    bl_label = "Build Dolly Camera Rig"
-    bl_description = "Build a Camera Dolly Rig"
-    bl_options = {'REGISTER', 'UNDO'}
 
-    def execute(self, context):
-        # build the Widgets
-        if not "WDGT_Camera_root" in bpy.context.scene.objects:
-            create_root_widget(self, "Camera_root")
-        if not "WDGT_Camera" in bpy.context.scene.objects:
-            create_camera_widget(self, "Camera")
-        if not "WDGT_Aim" in bpy.context.scene.objects:
-            create_aim_widget(self, "Aim")
-
-        # call the function to build the rig
-        build_dolly_rig(context)
-
-        return {'FINISHED'}
-
-
-def build_crane_rig(context):
-    # Define some useful variables:
-    boneLayer = (False, True, False, False, False, False, False, False,
-                 False, False, False, False, False, False, False, False,
-                 False, False, False, False, False, False, False, False,
-                 False, False, False, False, False, False, False, False)
-
-    # Add the new armature object:
-    bpy.ops.object.armature_add()
-    rig = context.active_object
-
-    # it will try to name the rig "Crane_rig" but if that name exists it will
-    # add .000 to the name
-    if "Crane_rig" not in context.scene.objects:
-        rig.name = "Crane_rig"
-    else:
-        rig.name = "Crane_rig.000"
-    rig["rig_id"] = "Crane_rig"
-
-    bpy.ops.object.mode_set(mode='EDIT')
-
-    # Remove default bone:
+def create_crane_bones(rig, bone_layers):
     bones = rig.data.edit_bones
-    bones.remove(bones[0])
 
-    # Add new bones:
+    # Add new bones
     root = bones.new("Root")
-    root.tail = (0.0, 3.0, 0.0)
+    root.tail = (0.0, 1.0, 0.0)
+    root.show_wire = True
 
-    ctrlAimChild = bones.new("aim_MCH")
-    ctrlAimChild.head = (0.0, 10.0, 1.0)
-    ctrlAimChild.tail = (0.0, 12.0, 1.0)
-    ctrlAimChild.layers = boneLayer
+    ctrl_aim_child = bones.new("Aim_shape_rotation-MCH")
+    ctrl_aim_child.head = (0.0, 10.0, 1.7)
+    ctrl_aim_child.tail = (0.0, 11.0, 1.7)
+    ctrl_aim_child.layers = bone_layers
 
-    ctrlAim = bones.new("Aim")
-    ctrlAim.head = (0.0, 10.0, 1.0)
-    ctrlAim.tail = (0.0, 12.0, 1.0)
+    ctrl_aim = bones.new("Aim")
+    ctrl_aim.head = (0.0, 10.0, 1.7)
+    ctrl_aim.tail = (0.0, 11.0, 1.7)
+    ctrl_aim.show_wire = True
 
     ctrl = bones.new("Camera")
-    ctrl.head = (0.0, 1.0, 1.0)
-    ctrl.tail = (0.0, 3.0, 1.0)
+    ctrl.head = (0.0, 1.0, 1.7)
+    ctrl.tail = (0.0, 2.0, 1.7)
 
     arm = bones.new("Crane_arm")
-    arm.head = (0.0, 0.0, 1.0)
-    arm.tail = (0.0, 1.0, 1.0)
+    arm.head = (0.0, 0.0, 1.7)
+    arm.tail = (0.0, 1.0, 1.7)
 
     height = bones.new("Crane_height")
     height.head = (0.0, 0.0, 0.0)
-    height.tail = (0.0, 0.0, 1.0)
+    height.tail = (0.0, 0.0, 1.7)
 
-    # Setup hierarchy:
+    # Setup hierarchy
     ctrl.parent = arm
     ctrl.use_inherit_rotation = False
     ctrl.use_inherit_scale = False
+    ctrl.show_wire = True
 
     arm.parent = height
     arm.use_inherit_scale = False
 
     height.parent = root
-    ctrlAim.parent = root
-    ctrlAimChild.parent = ctrlAim
-
-    # change display to BBone: it just looks nicer
-    bpy.context.object.data.display_type = 'BBONE'
-    # change display to wire for object
-    bpy.context.object.display_type = 'WIRE'
+    ctrl_aim.parent = root
+    ctrl_aim_child.parent = ctrl_aim
 
-    # jump into pose mode and change bones to euler
-    bpy.ops.object.mode_set(mode='POSE')
-    for x in bpy.context.object.pose.bones:
-        x.rotation_mode = 'XYZ'
-
-    # lock the relevant loc, rot and scale
-    bpy.context.object.pose.bones[
-        "Crane_arm"].lock_rotation = [False, True, False]
-    bpy.context.object.pose.bones["Crane_arm"].lock_scale = [True, False, True]
-    bpy.context.object.pose.bones["Crane_height"].lock_location = [True, True, True]
-    bpy.context.object.pose.bones["Crane_height"].lock_rotation = [True, True, True]
-    bpy.context.object.pose.bones["Crane_height"].lock_scale = [True, False, True]
-
-    # add the custom bone shapes
-    bpy.context.object.pose.bones["Root"].custom_shape = bpy.data.objects[
-        "WDGT_Camera_root"]  # add the widget as custom shape
-    # set the wireframe checkbox to true
-    bpy.context.object.data.bones["Root"].show_wire = True
-    bpy.context.object.pose.bones[
-        "Aim"].custom_shape = bpy.data.objects["WDGT_Aim"]
-    bpy.context.object.data.bones["Aim"].show_wire = True
-    bpy.context.object.pose.bones["Aim"].custom_shape_transform = bpy.data.objects[
-        rig.name].pose.bones["aim_MCH"]  # sets the "At" field to the child
-    bpy.context.object.pose.bones[
-        "Camera"].custom_shape = bpy.data.objects["WDGT_Camera"]
-    bpy.context.object.data.bones["Camera"].show_wire = True
-
-    # jump into object mode
+    # Jump into object mode
     bpy.ops.object.mode_set(mode='OBJECT')
+    pose_bones = rig.pose.bones
+
+    # Lock the relevant loc, rot and scale
+    pose_bones["Crane_arm"].lock_rotation = (False, True, False)
+    pose_bones["Crane_arm"].lock_scale = (True, False, True)
+    pose_bones["Crane_height"].lock_location = (True, True, True)
+    pose_bones["Crane_height"].lock_rotation = (True, True, True)
+    pose_bones["Crane_height"].lock_scale = (True, False, True)
+
+
+def build_camera_rig(context, mode):
+    bone_layers = tuple(i == 1 for i in range(32))
+    view_layer = bpy.context.view_layer
 
-    # Add constraints to bones:
-    con = rig.pose.bones['aim_MCH'].constraints.new('COPY_ROTATION')
+    rig_name = mode.capitalize() + "_Rig"
+    rig_data = bpy.data.armatures.new(rig_name)
+    rig = object_utils.object_data_add(context, rig_data, name=rig_name)
+    rig["rig_id"] = "%s" % rig_name
+    view_layer.objects.active = rig
+    rig.location = context.scene.cursor.location
+
+    bpy.ops.object.mode_set(mode='EDIT')
+
+    # Add new bones
+    if mode == "DOLLY":
+        create_dolly_bones(rig, bone_layers)
+    elif mode == "CRANE":
+        create_crane_bones(rig, bone_layers)
+
+    # Jump into object mode and change bones to euler
+    bpy.ops.object.mode_set(mode='OBJECT')
+    pose_bones = rig.pose.bones
+    for b in pose_bones:
+        b.rotation_mode = 'XYZ'
+
+    # Add custom properties to the armature’s Camera bone,
+    # so that all properties may be animated in a single action
+    # Add driver after the camera is created
+
+    # Lens property
+    pb = pose_bones['Camera']
+    pb["lens"] = 50.0
+    prop = rna_idprop_ui_prop_get(pb, "lens", create=True)
+    prop["default"] = 50.0
+    prop["min"] = 1.0
+    prop["max"] = 1000000.0
+    prop["soft_max"] = 5000.0
+
+    # DOF Focus Distance property
+    pb = pose_bones['Camera']
+    pb["focus_distance"] = 10.0
+    prop = rna_idprop_ui_prop_get(pb, "focus_distance", create=True)
+    prop["default"] = 10.0
+    prop["min"] = 0.0
+
+    # DOF F-Stop property
+    pb = pose_bones['Camera']
+    pb["aperture_fstop"] = 2.8
+    prop = rna_idprop_ui_prop_get(pb, "aperture_fstop", create=True)
+    prop["default"] = 2.8
+    prop["min"] = 0.0
+    prop["soft_min"] = 0.1
+    prop["soft_max"] = 128.0
+
+    # Build the widgets
+    root_widget = create_root_widget("Camera_Root")
+    camera_widget = create_camera_widget("Camera")
+    aim_widget = create_aim_widget("Aim")
+
+    # Add the custom bone shapes
+    pose_bones["Root"].custom_shape = root_widget
+    pose_bones["Aim"].custom_shape = aim_widget
+    pose_bones["Camera"].custom_shape = camera_widget
+
+    # Set the "At" field to the child
+    pose_bones["Aim"].custom_shape_transform = pose_bones["Aim_shape_rotation-MCH"]
+
+    # Add constraints to bones
+    con = pose_bones['Aim_shape_rotation-MCH'].constraints.new('COPY_ROTATION')
     con.target = rig
     con.subtarget = "Camera"
 
-    con = rig.pose.bones['Camera'].constraints.new('TRACK_TO')
+    con = pose_bones['Camera'].constraints.new('TRACK_TO')
     con.target = rig
     con.subtarget = "Aim"
     con.use_target_z = True
 
-    # Add custom Bone property to Camera bone
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "lock", create=True)
-    ob["lock"] = 1.0
-    prop["soft_min"] = prop["min"] = 0.0
-    prop["soft_max"] = prop["max"] = 1.0
-
-    # Add Driver to Lock/Unlock Camera from Aim Target
-    rig = bpy.context.view_layer.objects.active
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-
-    constraint = pose_bone.constraints["Track To"]
-    inf_driver = constraint.driver_add('influence')
-    inf_driver.driver.type = 'SCRIPTED'
-    var = inf_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["lock"]'
-    inf_driver.driver.expression = 'var'
-
-    # Add custom property for the lens / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "focal_length", create=True)
-    ob["focal_length"] = 50.0
-    prop["soft_min"] = prop["min"] = 1.0
-    prop["default"] = 50.0
-    prop["soft_max"] = prop["max"] = 5000.0
-
-    # Add custom property for the focus distance / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "focus_distance", create=True)
-    ob["focus_distance"] = 10.00
-    prop["soft_min"] = prop["min"] = 0.0
-    prop["soft_max"] = prop["max"] = 1000.0
-
-    # Add custom property for the focus distance / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "focus_distance", create=True)
-    ob["focus_distance"] = 10.00
-    prop["soft_min"] = prop["min"] = 0.0
-    prop["soft_max"] = prop["max"] = 1000.0
-
-    # Add custom property for the f-stop / add the driver after the camera is created
-    ob = bpy.context.object.pose.bones['Camera']
-    prop = rna_idprop_ui_prop_get(ob, "f-stop", create=True)
-    ob["f-stop"] = 2.8
-    prop["soft_min"] = prop["min"] = 0.1
-    prop["soft_max"] = prop["max"] = 128.00
-
-    # Add the camera object:
-    bpy.ops.object.mode_set(mode='OBJECT')
+    # Change display to BBone: it just looks nicer
+    bpy.context.object.data.display_type = 'BBONE'
+    # Change display to wire for object
+    bpy.context.object.display_type = 'WIRE'
+
+    # Add the camera object
+    cam_name = "%s_Camera" % mode.capitalize()
+    cam_data = bpy.data.cameras.new(cam_name)
+    cam = object_utils.object_data_add(context, cam_data, name=cam_name)
+    view_layer.objects.active = cam
+    context.scene.camera = cam
+
+    cam.data.display_size = 1.0
+    cam.rotation_euler[0] = pi / 2.0  # Rotate the camera 90 degrees in x
 
-    bpy.ops.object.camera_add()
-    cam = bpy.context.active_object
-
-    # this will name the Camera Object
-    if 'Crane_camera' not in context.scene.objects:
-        cam.name = "Crane_camera"
-    else:
-        cam.name = "Crane_camera.000"
-
-    # this will name the camera Data Object
-    if "Crane_camera" not in bpy.context.scene.objects.data.camera:
-        cam.data.name = "Crane_camera"
-    else:
-        cam.data.name = "Crane_camera.000"
-
-    cam_data_name = bpy.context.object.data.name
-    bpy.data.cameras[cam_data_name].display_size = 1.0
-    cam.rotation_euler = [radians(90), 0, 0]  # rotate the camera 90 degrees in x
-    cam.location = (0.0, -2.0, 0.0)  # move the camera to the correct postion
+    cam.location = (0.0, -1.0, 0.0)  # Move the camera to the correct position
     cam.parent = rig
     cam.parent_type = "BONE"
     cam.parent_bone = "Camera"
 
-    # Add Driver to link the camera lens to the custom property on the armature
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-    lens_driver = cam.data.driver_add("lens")
-    lens_driver.driver.type = 'SCRIPTED'
-    var = lens_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["focal_length"]'
-    lens_driver.driver.expression = 'var'
-
-    # Add Driver to link the camera focus distance to the custom property on the armature
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-    lens_driver = cam.data.driver_add("dof.focus_distance")
-    lens_driver.driver.type = 'SCRIPTED'
-    var = lens_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["focus_distance"]'
-    lens_driver.driver.expression = 'var'
-
-    # Add Driver to link the camera f-stop to the custom property on the armature
-    pose_bone = bpy.data.objects[rig.name].pose.bones['Camera']
-    lens_driver = cam.data.driver_add("dof.aperture_fstop")
-    lens_driver.driver.type = 'SCRIPTED'
-    var = lens_driver.driver.variables.new()
-    var.name = 'var'
-    var.type = 'SINGLE_PROP'
-
-    # Target the Custom bone property
-    var.targets[0].id = bpy.data.objects[rig.name]
-    var.targets[0].data_path = 'pose.bones["Camera"]["f-stop"]'
-    lens_driver.driver.expression = 'var'
-
-    #  the location/rotation/scale of the camera
-    cam.lock_location = [True, True, True]
-    cam.lock_rotation = [True, True, True]
-    cam.lock_scale = [True, True, True]
-
-    # Set new camera as active camera
-    bpy.context.scene.camera = cam
-
-    # make sure the camera is selectable by default (this can be locked in the UI)
-    bpy.context.object.hide_select = False
-
-    # make the rig the active object before finishing
-    bpy.context.view_layer.objects.active = rig
-    bpy.data.objects[cam.name].select_set(False)
-    bpy.data.objects[rig.name].select_set(True)
+    # Lock camera transforms
+    cam.lock_location = (True,) * 3
+    cam.lock_rotation = (True,) * 3
+    cam.lock_scale = (True,) * 3
+
+    # Add drivers to link the camera properties to the custom props
+    # on the armature
+    for prop_from, prop_to in (("lens", "lens"),
+                               ("focus_distance", "dof.focus_distance"),
+                               ("aperture_fstop", "dof.aperture_fstop")):
+        driver = cam.data.driver_add(prop_to)
+        driver.driver.type = 'SCRIPTED'
+        var = driver.driver.variables.new()
+        var.name = 'var'
+        var.type = 'SINGLE_PROP'
+
+        # Target the custom bone property
+        var.targets[0].id = rig
+        var.targets[0].data_path = 'pose.bones["Camera"]["%s"]' % prop_from
+        driver.driver.expression = 'var'
+
+    # Make the rig the active object
+    for ob in view_layer.objects:
+        ob.select_set(False)
+    rig.select_set(True)
+    view_layer.objects.active = rig
 
     return rig
 
 
-class ADD_CAMERA_RIGS_OT_build_crane_rig(Operator):
-    bl_idname = "add_camera_rigs.build_crane_rig"
-    bl_label = "Build Crane Camera Rig"
-    bl_description = "Build a Camera Crane Rig"
+class OBJECT_OT_build_camera_rig(Operator):
+    bl_idname = "object.build_camera_rig"
+    bl_label = "Build Camera Rig"
+    bl_description = "Build a Camera Rig"
     bl_options = {'REGISTER', 'UNDO'}
 
-    def execute(self, context):
-        # build the Widgets
-        if not "WDGT_Camera_root" in bpy.context.scene.objects:
-            create_root_widget(self, "Camera_root")
-        if not "WDGT_Camera" in bpy.context.scene.objects:
-            create_camera_widget(self, "Camera")
-        if not "WDGT_Aim" in bpy.context.scene.objects:
-            create_aim_widget(self, "Aim")
-
-        # call the function to build the rig
-        build_crane_rig(context)
+    mode: bpy.props.EnumProperty(items=
+                                 (('DOLLY',) * 3,
+                                  ('CRANE',) * 3,),
+                                 name="mode",
+                                 description="", default="DOLLY")
 
+    def execute(self, context):
+        # Build the rig
+        build_camera_rig(context, self.mode)
         return {'FINISHED'}
 
 
-# dolly and crane entries in the Add Object > Camera Menu
 def add_dolly_crane_buttons(self, context):
+    """Dolly and crane entries in the Add Object > Camera Menu"""
     if context.mode == 'OBJECT':
-        self.layout.operator(
-            ADD_CAMERA_RIGS_OT_build_dolly_rig.bl_idname,
+        op = self.layout.operator(
+            OBJECT_OT_build_camera_rig.bl_idname,
             text="Dolly Camera Rig",
             icon='CAMERA_DATA'
         )
-        self.layout.operator(
-            ADD_CAMERA_RIGS_OT_build_crane_rig.bl_idname,
+        op.mode = "DOLLY"
+        op = self.layout.operator(
+            OBJECT_OT_build_camera_rig.bl_idname,
             text="Crane Camera Rig",
             icon='CAMERA_DATA'
         )
+        op.mode = "CRANE"
 
 
 classes = (
-    ADD_CAMERA_RIGS_OT_build_dolly_rig,
-    ADD_CAMERA_RIGS_OT_build_crane_rig,
+    OBJECT_OT_build_camera_rig,
 )