*args, **kwargs, and custom class wrappers in Python

At Thyme Care, we’re growing quickly, so we are always looking for ways to improve the onboarding process for new developers and to make application changes easier. One way we do this is by abstracting out the cookie-cutter code written for similar features so there is less cognitive load in making changes. Recently, we saw an opportunity for this in our data models and API interactions.

The Backstory: Adding custom attributes to SQLAlchemy columns

One of the places where we started to see a lot of repeated logic in our codebase was around adding a new data component to the RESTful portion of our API. The common pattern when adding a new data table was:

• Update the ORM: Create a new SQLAlchemy data model using the declarative approach

• Add a migration: Auto-generate a new alembic migration to add the model to the database

• Add data validation: Create a mostly auto-generated marshmallow schema for the model to power our validation and Swagger API docs

• Add endpoints: Create endpoints for getting, updating, and deleting the new data

• Add mock data: Add a new factory to generate mock data for the data model using factory_boy and faker

• Write tests: Create tests to guarantee the API layer business logic

From a business logic side, there are some columns that can be set only on create, some columns that can be set on create and update, and some columns that are required/optional for one/both of those operations. SQLAlchemy doesn’t provide us a way to specify this on the column, so we were redefining this information at our API layer:

UpdateAccountSchema = Schema.from_dict(
   {
       "some_field": fields.Str(
           validate=validate.OneOf([x.value for x in Accounts.SomeField])
       ),
       "other_field": fields.Str(
           validate=validate.OneOf([x.value for x in Accounts.OtherField])
       ),
   }
)

We were able to pull out much of this into a function to read from the SQLAlchemy model, so eventually got to this abstraction to build our marshmallow schema:

UpdateAccountSchema = simple_auto_schema(
   Account,
   name="UpdateAccountSchema",
   optional_fields=[
       "some_field",
       "other_field",
   ],
)

This was better but still felt like we were duplicating business logic in multiple places. For instance, if a field is nullable in the SQLAlchemy model, it should be optional at the API layer in most cases. This meant when adding a new field you had to add it not only to the SQLAlchemy model but to a bunch of scattered schemas to get it to be accepted at the API layer. 

Ideally, we wanted to put all of the information about a column inline with the column declaration on the SQLAlchemy model. It’s not a foolproof solution for more complicated models or endpoints but for your basic CRUD endpoints it’s a nice thing to opt in to. What we wanted was this:

first_name = Column(String, nullable=False, user_editable=True)

With that code, our create and update endpoints can know to accept a required first_name field. 

A basic multiple inheritance pattern doesn’t work here, since the SQLAlchemy column constructor doesn’t allow extra arguments.

# the error passing an extra argument to the SQLAlchemy constructor
SAWarning: Can't validate argument 'user_creatable'; can't locate any SQLAlchemy dialect named 'user'
super().__init__(*args, **kwargs, user_creatable=user_creatable)

To do make this work, we wrote a wrapper around the base SQLAlchemy column class that we use instead, which injects our own metadata:

class BaseColumn(_SQLAlchemyColumn):
   def __init__(
       self,
       *args,
       user_creatable: Optional[bool]=None,
       user_editable: bool=False,
       user_serialize: bool=True,
       **kwargs,
   ):
       """
       Wrapper around SQLAlchemy `Column` that provides some column metadata that
       can be used to automate parts of endpoint schemas and `to_dict` dumping

       :param args: args to pass to SQLAlchemy Column creation
       :param user_creatable: Whether an app user can alter a value for this column
       on item creation. If not provided, this takes the value of `user_editable`
       :param user_editable: Whether an app user can alter a value for this column
       through an update endpoint. Defaults to False
       :param user_serialize: Whether to include this column in serialization 
       Defaults to True
       :param kwargs: kwargs to pass to SQLAlchemy Column creation
       """
       self.user_serialize = user_serialize

       if user_creatable is not None:
           self.user_creatable = user_creatable
       else:
           self.user_creatable = user_editable
       self.user_editable = user_editable
       super().__init__(*args, **kwargs)

Finally, we can put them inline:

some_immutable_var = BaseColumn(String, nullable=False, user_creatable=True)
some_editable_var = BaseColumn(String, user_editable=True)
some_secret_var = BaseColumn(String, user_creatable=True, user_serializable=False)

Then, in generating our schemas for our endpoints, we can use some simple methods we’ve added to our tables (full code for BaseModel in the appendix):

CreatePersonSchema = simple_auto_schema(
   Person,
   name="CreatePersonSchema",
   required_fields=BaseModel.user_creatable_cols(Person, nullable=False),
   optional_fields=BaseModel.user_creatable_cols(Person, nullable=True),
)

One of the downsides of this is with documentation in IDEs in that we lose the docstring and argument hints for the underlying class because it’s overridden by our wrapper. This isn’t a huge problem if the underlying class is a well-known class with extensive documentation (like the SQLAlchemy Column below) but is a consideration.

All in all, we’ve had luck using this wrapper pattern to make our code more maintainable and thought it’d be helpful to share out. We’d love to hear any thoughts on more advantages, downsides, or other ideas in the comments, and we’re looking forward to sharing more learnings in the future.

tl;dr; Generic Example

Sometimes we want to add some custom data to a third-party library class in Python but obviously can’t change the library itself. Using a simple inheritance pattern along with Python’s *args and **kwargs symbols, we can insert our own metadata into a wrapper class without affecting the underlying implementation. Here’s the basic pattern using a pretend Square class library:

# The library code with an __init__ function that has some complicated set of arguments
class WeirdClass(object):
   def __init__(self, first_arg, *args, some_arg=1, other_arg=2):
      ...

# Our code
from some.library import WeirdClass as _SomeLibraryWeirdClass
class BaseWeirdClass(_SomeLibraryWeirdClass):
   def __init__(
       self,
       *args,
       editable=False, # <- Our custom class attributes
       **kwargs,
   ):
       """
       Wrapper around SomeLibrary’s `WeirdClass` that adds whether it’s editable
as an attribute

       :param args: args to pass to SomeLibrary’s WeirdClass creation
       :param editable: Whether our app logic should let the user can change the weird class instances
       :param kwargs: kwargs to pass to SomeLibrary’s WeirdClass creation
       """
       self.editable = editable
       super().__init__(*args, **kwargs) # **kwargs won’t contain `editable`, so our wrapper is invisible to the underlying implementation

Usage:
BaseWeirdClass(2, 4, 6, 8, some_arg=2, editable=True, other_arg=3)

Appendix: Underlying helper code used in examples above

from marshmallow_SQLAlchemy import SQLAlchemyAutoSchema, SQLAlchemySchema, auto_field

def dynamic_meta(model, meta_opts=None):
   meta_opts = meta_opts or {}
   return type("Meta", (BaseMeta,), {"model": model, **meta_opts})


def simple_auto_schema(
   model,
   name=None,
   required_fields=None,
   optional_fields=None,
   meta_opts=None,
):
   """
   :param model: SQLAlchemy model to create a schema for
   :param name: name of the schema, if not the barebones schema for the model
   :param required_fields: list of required field names to use instead of all fields
   on the model
   :param optional_fields: list of optional field names to use instead of all fields
   on the model
   :param meta_opts: options to pass to the SQLAlchemy Meta class

   :return: marshmallow Schema class
   """
   props = {}
   base_classes = [SQLAlchemyAutoSchema]
   if required_fields or optional_fields:
       base_classes = [SQLAlchemySchema]
       assert name is not None, "unique name required if explicit fields are given"
       for f in optional_fields or []:
           props[f] = auto_field(required=False)
       for f in required_fields or []:
           props[f] = auto_field(required=True)

   # Dynamically create a new schema type that inherits from base_classes
   return type(
       name or (model.__name__ + "Schema"),
       tuple(base_classes),
       {
         # Dynamically create a SQLAlchemy Meta class with the 
# given model and any provided options
           "Meta": dynamic_meta(model, meta_opts),
       }
       | props,
   )

class BaseModel(db.Model):
   __abstract__ = True

   @staticmethod
   def cols_to_serialize(model):
       """
       Gets a list of column names that should be included in item serialization

       :param model: Table class or instance of SQLAlchemy object
       :return: List of column names
       """
       return [c.name for c in model.__table__.columns if c.user_serialize]

   @staticmethod
   def _filter_cols(cols, nullable):
       if nullable is True:
           return [c.name for c in cols if c.nullable]
       elif nullable is False:
           return [c.name for c in cols if not c.nullable]
       return [c.name for c in cols]

   @staticmethod
   def user_creatable_cols(model, nullable=None):
       """
       Gets a list of column names that should be alterable on item creation
       by an app user.

       :param model: Table class or instance of SQLAlchemy object
       :param nullable: Boolean whether to filter for columns that are only nullable,
       or only non-nullable. If `None`, no filtering is applied
       :return: List of column names
       """
       cols = [c for c in model.__table__.columns if c.user_creatable]
       return BaseModel._filter_cols(cols, nullable)

   @staticmethod
   def user_editable_cols(model, nullable=None):
       """
       Gets a list of column names that should be alterable on item edit
       by an app user.

       :param model: Table class or instance of SQLAlchemy object
       :param nullable: Boolean whether to filter for columns that are only nullable,
       or only non-nullable. If `None`, no filtering is applied
       :return: List of column names
       """
       cols = [c for c in model.__table__.columns if c.user_editable]
       return BaseModel._filter_cols(cols, nullable)