Adding to Unity’s Built-In Classes Using Extension Methods

(This article was originally published at third-helix.com.)

Have you ever found yourself wishing a built-in Unity class had some functionality that isn’t there? C# extension methods are the answer!

In this article, I’ll teach you how to use extension methods to add functionality to existing classes, no matter if they’re built-in Unity types, types defined in a third-party plugin, or even types defined in an Asset Store package which you could edit but you’re (rightly) worried about later package updates stomping your “patch”.

Seemingly obvious API omissions can be frustrating, but extension methods let you “fix” just about any API to your liking.

When Extension Methods Are Useful

Imagine you need a way to set the layer of a GameObject and all its children, but there’s no GameObject.SetLayerRecursively() available. You could embed a loop in your code:

// Do some work.

// Set the layer of this GameObject and all its children.
gameObject.layer = someLayer;
foreach(Transform t in transform)
    t.gameObject.layer = someLayer;

// Do some more work.

This will work fine, although it’s not the cleanest thing in the world, and you’ll need to copy these lines of code around to every place where you want to do this operation. It would be better to encapsulate that code in a function, and make that function available to everyone via a “helper” class:

public class GameObjectHelper
{
    public static void SetLayerRecursively(GameObject gameObject, int layer)
    {
        gameObject.layer = someLayer;
        foreach(Transform t in transform)
            t.gameObject.layer = someLayer;
    }
}

public class Test : MonoBehaviour
{
    void Start()
    {
        GameObjectHelper.SetLayerRecursively(gameObject, someLayer);
    }
}

This works fine too, but it can be easy to forget to use the helper class to invoke a function that feels like it belongs in GameObject itself. Wouldn’t it be nicer if you could just do this?

gameObject.SetLayerRecursively(someLayer);

As it turns out, C# lets you make this happen!

Declaring An Extension Method

Extension methods are declared just like regular functions, except that a) they must be static, b) they must be declared inside a non-generic, non-nested, static class, and c) they include a special syntax for their first argument, which is a reference to the calling object:

public static class GameObjectExtensions
{
    public static void SetLayerRecursively(this GameObject gameObject, int layer)
    {
        // Implementation goes here
    }
}

Though declared as static, this is invoked as if it were an instance method:

myGameObject.SetLayerRecursively(someLayer);

Note that when we call the method we actually omit the first argument and skip straight to the second. Take another look at the method declaration above. See how the first argument is declared using the “this” keyword? That’s what tells the compiler to infer that argument as the calling object; in this case, myGameObject.

That’s actually all there is to it. Extension methods are easy!

For what it’s worth, I like to organize my extension methods into classes named ClassNameExtensions. So I have GameObjectExtensions, TransformExtensions, ColorExtensions, and so on. There’s nothing that says you have to do this; I just like the organization. You could pack them all together into a single Extensions class if you prefer, or any other kind of organization you want. Just remember that extensions methods must be declared inside a non-generic, non-nested, static class; the compiler will complain otherwise.

Limitations of Extension Methods

Extension methods can add to an existing class, but they cannot override existing methods. For example, if you declared this:

public static class GameObjectExtensions
{
    public static Component GetComponent(this GameObject gameObject, Type type)
    {
        Debug.LogError(“LOL, you got trolled”);
    }
}

…the compiler would basically ignore you. The rule for multiple method declarations using the same signature is: instance methods always take precedence over extension methods.

To extend a type, you’ll need make sure it’s in scope with the “using” directive. If you’re extending any built-in Unity type, you’re covered by virtue of the “using UnityEngine” that’s a standard entry in most Unity scripts. If you’re extending an editor type, you’ll need to add “using UnityEditor”, just like you would if you were calling that type. If you’re extending a type from a third-party plugin, you may need to import a namespace; check the plugin’s documentation (or source code, if you have it) for details.

I’ve already mentioned that extension methods must be declared inside a non-generic, non-nested, static class, which means you can’t just drop them in willy-nilly wherever you want. In practice, this “limitation” turns out to be a useful organizing device. You can certainly argue that this is virtually identical to the “helper class” example at the top of this article, and in terms of implementation effort that’s probably true; the difference is that with extension methods you get a cleaner, more normalized calling syntax. This really comes down to personal preference.

Some Useful Extension Methods

For the rest of this article, I’ll share some useful extension methods I’ve written over the past few months. Feel free to incorporate these into your own codebase and modify them at-will. (But don’t just copy-paste them into a text file and try to sell it on the Asset Store; that would make you a horrible person.)

Set Layer Recursively

I use the excellent 2D Toolkit for UI, and I use a two-camera setup: one camera for the scene, one camera for the UI. In order to make the UI camera render only UI objects, all the UI objects need to be on a UI layer. When I create UI objects from script I need to set that layer, and often I’m creating objects with children (like a button with a child sprite and a child text label).

It’d be nice to have a way to set the layer for this entire hierarchy with a single function call (you’ll recognize this as the example at the top of this article). Here’s what the call looks like:

myButton.gameObject.SetLayerRecursively(LayerMask.NameToLayer(“UI”));

And here’s that extension method:

// Set the layer of this GameObject and all of its children.
public static void SetLayerRecursively(this GameObject gameObject, int layer)
{
    gameObject.layer = layer;
    foreach(Transform t in gameObject.transform)
        t.gameObject.SetLayerRecursively(layer);
}

Set Visual/Collision Recursively

While we’re doing things recursively, wouldn’t it be nice if we could enable/disable just the renderers, or just the colliders, for an entire hierarchy? This can be useful for UI, but it can also be useful in the scene. Imagine a complex hierarchy that defines a really fancy animated force field, and a game mechanic whereby you can switch your avatar’s polarity, allowing passage through force fields of a particular color. When you switch polarity, you’d loop through the matching force fields and disable their colliders, like this:

foreach(ForceField forceField in forceFields)
    forceField.gameObject.SetCollisionRecursively(false);

Here’s that extension method:

public static void SetCollisionRecursively(this GameObject gameObject, bool tf)
{
    Collider[] colliders = gameObject.GetComponentsInChildren();
    foreach(Collider collider in colliders)
        collider.enabled = tf;
}

And the same principle applies for renderers:

public static void SetVisualRecursively(this GameObject gameObject, bool tf)
{
    Renderer[] renderers = gameObject.GetComponentsInChildren();
    foreach(Renderer renderer in renderers)
        renderer.enabled = tf;
}

Filter Child Components By Tag

It’s easy to search for child components using GameObject.GetComponentsInChildren(), but what if you have a have hierarchy in which you have lots of instances of a type, and you only want those instances which have a particular tag? In my case, I had a compound object with numerous renderers, and I needed to drive a material color on a subset of those renderers tagged as “Shield”.

This proved convenient:

m_renderers = gameObject.GetComponentsInChildrenWithTag("Shield");

Here’s that extension method:

public static T[] GetComponentsInChildrenWithTag(this GameObject gameObject, string tag)
    where T: Component
{
    List results = new List();

    if(gameObject.CompareTag(tag))
        results.Add(gameObject.GetComponent());

    foreach(Transform t in gameObject.transform)
        results.AddRange(t.gameObject.GetComponentsInChildrenWithTag(tag));

    return results.ToArray();
}

I noted earlier that extension methods must be declared inside non-generic classes. That doesn’t mean the extension method itself can’t be generic, however! When we call this method we replace T with the type we’re interested in — in the preceding call example, it was Renderer — and that type is used for each occurrence of T in the method implementation. The “where” keyword specifies that T must be of type Component or a type derived from Component (the compiler will throw an error otherwise).

See this MSDN article for more information about generics.

Get Components In Parents

It’s all well and good to get components in children, but sometimes you need to search up. A common case I run into is when figuring out what to do with a collision result. I have a Player type, and its hierarchy is made up of several GameObjects which represent visuals, colliders, equipped items, and so on. Typically when I get a collision result on a player, the collider is bound to a child GameObject, so I can’t just do GetComponent() and then player.TakeDamage() or whatever, because there’s no Player component on the GameObject I actually hit. In this case I need to search up the hierarchy and find the Player to which this collider is parented; also the Player may not necessarily be this collider’s immediate parent.

So now I do this:

Player player = collider.gameObject.GetComponentInParents();

Here’s that extension method:

public static T GetComponentInParents(this GameObject gameObject)
    where T : Component
{
    for(Transform t = gameObject.transform; t != null; t = t.parent)
    {
        T result = t.GetComponent();
        if(result != null)
            return result;
    }

    return null;
}

Just like Unity has both GameObject.GetComponentInChildren (singular) and GameObject.GetComponentsInChildren (plural), I also created a version that gets all components in parents:

public static T[] GetComponentsInParents(this GameObject gameObject)
    where T: Component
{
    List results = new List();
    for(Transform t = gameObject.transform; t != null; t = t.parent)
    {
        T result = t.GetComponent();
        if(result != null)
            results.Add(result);
    }

    return results.ToArray();
}

Note: it would be trivial to create a GetComponentsInParentsWithTag, but I haven’t run into a need for it yet. If you’d like to exercise your newfound knowledge of extension methods, this might be a good exercise.

Get An Object’s Collision Mask

Here’s one whose omission is particularly perplexing. You can get the layer a GameObject is on, which is useful for both rendering and collision purposes, but there’s no easy way to figure out the set of layers that GameObject can collide against.

I ran into this when implementing weapons. Projectile-based weapons are simple: they get a Rigidbody and a Collider of some sort, and the collision system handles everything for me. But a rail gun-like weapon is a different story: there’s no Rigidbody, just a script-invoked raytest. You can pass a collision mask — a bitfield — into a raytest, but what if you want the collision mask to be based on the weapon’s layer? It’d be nice to set some weapons to “Team1” and others to “Team2”, perhaps, and also to ensure your code doesn’t break if you change the collision matrix in the project’s Physics Settings.

Really, I wanted to just do this:

if(Physics.Raycast(startPosition, direction, out hitInfo, distance,
    weapon.gameObject.GetCollisionMask())
)
{
    // Handle a hit
}

That raycast will only hit objects which the calling weapon is allowed to collide with, based on its layer and the project’s collision matrix.

Here’s that extension method:

public static int GetCollisionMask(this GameObject gameObject, int layer = -1)
{
    if(layer == -1)
        layer = gameObject.layer;

    int mask = 0;
    for(int i = 0; i < 32; i++)
        mask |= (Physics.GetIgnoreLayerCollision(layer, i) ? 0 : 1) << i;

    return mask;
}

Note the optional “layer” argument. If omitted, it uses the layer of the calling GameObject, which is the most common/intuitive case (for me, at least). But you can specify a layer and it’ll hand you the collision mask for that layer instead.

Easily Change A Color’s Alpha

I often find myself wanting to modulate the alpha value of a color without changing the color itself, for blinking/pulsing effects. Because Color is a struct, and structs in C# are immutable, you can’t simply assign color.a; you’ll get a compiler error. Fortunately, extension methods can extend structs as well as classes:

public static Color WithAlpha(this Color color, float alpha)
{
    return new Color(color.r, color.g, color.b, alpha);
}

This method makes modulating a color’s alpha clean and simple:

GUI.color = desiredColor.WithAlpha(currentAlpha);

Performance Considerations

There really aren’t any. Extension methods are a compile-time device; once you reach runtime, they look and act just like any other method call. So it doesn’t matter whether you write an extension method or an explicit helper class; as long as the guts of the method are identical, both implementations should perform identically. You could theoretically squeeze a vanishingly tiny bit of performance out of simply inlining all the code, avoiding the overhead of any function call at all, but in this day and age that’s completely pointless to worry about unless you’re calling the function millions of times.

Extension methods are mainly useful as a way of cleaning up and normalizing your code. They have little or no effect on the behavior of that code.

Conclusion

I hope this article has been clear and useful. Extension methods are one of my favorite features of C#, because I’m a bit obsessive about having clean, easy-to-read code and they can really help make that happen. Of course, there’s more than one way to skin a cat, and I’m by no means suggesting that extension methods are the right or only way to do things. You can certainly argue that making an explicit helper class containing regular static functions — like GameObjectHelper.SetActiveRecursively() in the example at the top of this article — is just as good as an extension method-based implementation of the same; they’re six of one, half-dozen of another. I prefer the extension method approach because it feels like a natural and intuitive API extension, rather than a “bolt-on”, but that’s strictly a matter of personal preference.

Like every coding practice, extension methods are just one tool in the toolbox. But I encourage you to give them a shot!

(Josh Sutphin is an indie game developer, former lead designer of Starhawk (PS3), and creator of the Ludum Dare-winning RTS/tower-defense hybrid Fail-Deadly. He blogs at third-helix.com and tweets nonsense at @invicticide.)