Experimenting extending Shotgun to return custom objects instead of simple dicts.

shotgunWrapTest.py

import pprint
import sys
import shotgun_api3 as sg


class Shotgun(sg.Shotgun):
	__registry = {}

	def __init__(self, *args, **kwargs):
		super(Shotgun, self).__init__(*args, **kwargs)

	def find(self, *args, **kwargs):
		return [self.wrap(o) for o in super(Shotgun, self).find(*args, **kwargs)]

	def find_one(self, *args, **kwargs):
		return self.wrap(super(Shotgun, self).find_one(*args, **kwargs))

	def registerWrapper(self, entityType, wrapper):
		if not issubclass(wrapper, dict):
			raise TypeError('Wrapper should be a dict subclass.')

		self.__registry[entityType] = wrapper

	@classmethod
	def wrap(cls, obj):
		if obj is not None and 'type' in obj:
			wrapperClass = cls.__registry.get(obj['type'])
			if wrapperClass:
				wrapper = wrapperClass()
				wrapper.update(obj)
				return wrapper
		return obj


def getShotgunConnection(url, name, key, registry):
	con = Shotgun(url, name, key)
	for k, v in registry.iteritems():
		con.registerWrapper(k, v)
	return con


class ShotgunObjectWrapper(dict):
	def __getitem__(self, key):
		value = super(ShotgunObjectWrapper, self).__getitem__(key)
		return Shotgun.wrap(value)


class Shot(ShotgunObjectWrapper):
	pass


class Ticket(ShotgunObjectWrapper):
	pass


class Tool(ShotgunObjectWrapper):
	pass


def main(args):
	con = getShotgunConnection(
		'https://oblique.shotgunstudio.com',
		'$YOUR_NAME$',
		'$YOUR_KEY$',
		{
			'Shot': Shot,
			'Ticket': Ticket,
			'Tool': Tool,
		}
	)

	ticket = con.find_one('Ticket', [['id', 'is', 57]], ['id', 'title', 'sg_tool'])
	pprint.pprint(ticket)
	print type(ticket)
	pprint.pprint(ticket['sg_tool'])
	print type(ticket['sg_tool'])


if __name__ == '__main__':
	sys.exit(main(sys.argv))