Introduction
Strongly typed languages need some mechanism to express generic functions, other languages don't enforce types which means every function is generic. Many languages implement some form of this feature under a variety of different names:
Python
In dynamic languages like JavaScript and Python, you can pass any object to any function. But you must make sure the object implements the methods you call in that function, or you'll get a nasty runtime error:
Python3.8 introduced a typing feature named protocols which is related to traits, here's an example of a protocol named Shape that defines an abstract function signature area but doesn't implement it:
This is a similar concept to Mojo's traits in terms of syntax and concept, but it only gives linter warnings and some quality of life improvements like better completions. With Mojo, you also get no performance penalty for writing generic functions across different types, and you can't put mistakes into production as it simply won't compile.
Mojo 🔥
Let's take a look at Mojo traits:
We can now create a function that accepts anything implementing the Shape trait and run the abstract area method in the function:
The [T: Shape] at the start is common syntax across languages, you can think of T as declaring a generic Type that is constrained by the Shape trait.
Now we'll create a concrete fn area implementation on Circle:
Circle(Shape) means that the struct Circle must implement all the methods specified in the Shape trait, in this case it's just the fn area(self) -> Float64 signature.
You can now run it through the generic function:
If we remove fn area from Circle, the compiler won't allow us to build the program, so it's impossible to get a runtime error for this mistake. The other advantage in Mojo compared to Python, is the error tells us exactly what signature we need to implement:
Lets create another type that implements the Shape trait:
And run it through the same function:
Truly Zero-Cost Generics
The popular pattern from C++ to achieve this behavior was using inheritance and abstract classes, but the compiler can't reason about what types are used when running methods on an abstract class, which can have significant performance impacts. The pattern can also explode a code base in complexity.
C++ added multiple features to address these problems, such as templates and concepts. But they still aren't completely zero-cost, while Mojo can guarantee that values are register passable when using traits for truly zero-cost generics.
Multiple Traits
The __str__ method comes from Python, it determines what will happen when you print() the type. We added a Stringable trait in the standard library, which you can implement it on your type to make it printable:
Let's add the Stringable trait to Circle:
Circle now takes the Shape and Stringable traits, so it must implement fn area and fn __str__ to compile.
This allows us to print the type just like Python:
Trait Inheritance
A really cool feature of traits, is it allows users to compose their types with your library. We've added a few simple traits to our standard library so you can inherit them into your own types.
Here's an example of creating a new trait that inherits Shape, along with the standard library Stringable and CollectionElement so that you can push your type into a DynamicVector. The requirements for CollectionElement are implemented when using the @value decorator.
Now that we have a trait that's composed our three traits together, we can create a function that makes use of all of them:
Try defining your own Rectangle type that implements VecPrintableShape!
Database Trait
If you're still struggling to understand why traits are useful, a common example used to demonstrate the utility is a Database trait. We'll only define two methods to simplify the concept:
Now we can pass around an object that implements Database and use it's abstract methods:
For example, imagine you have a function that runs inference on an image and stores the result somewhere. Maybe we just want to use SQLite for our local batch tests, but in production it'll be stored in some dynamodb instance. Or maybe we just want the flexibility to change out the database later without causing breaking changes for users of our library.
Lets import the sqlite3 Python package to implement a database that runs locally, and conforms to the trait above:
Now we can pass the database to our previous function:
If you don't have sqlite installed, first run pip3 install sqlite3
Then run the program:
And now we can implement the Mongo version:
You'll need to pip install pymongo if you want to run this, and follow the instructions here to start a service.
Then run it:
This is simplified to demonstrate the functionality, but you could create an entire library following these principles wrapping Python libraries, and then introduce optimized Mojo implementations where you need better performance without changing the API.
Conclusion
You may have seen abstract methods navigating around Python code bases with ..., and not understood why they're there. It's pervasive in ML libraries where the authors want to provide correctness and nice tooling while still having multiple implementations for CUDA, CPU, and the many emerging hardware backends. You get an extra benefit in Mojo, you can write generic reusable functionality across multiple types, but still retain full type safety and compiler optimizations by writing concrete implementations for each type.
There are more features to come for traits such as default implementations, make sure to check back on the docs.
We're excited to see what you build with traits, please share your projects on the Discord and GitHub!
