Type polymorphic functions in TypeScript

Admittedly, the title might be bit broad. Polymorphism is a theoretical concept that’s deeply rooted in programming language theory, and it has many varieties. I am by no means an expert in programming language theory. So, I won’t use the term “polymorphic functions” in this blog post in a strict, academic sense. Rather, I will use it mainly to refer to functions in two ways: those that behave differently depending on their applied argument types (i.e. ad-hoc polymorphisms) and those that take a variable number of arguments (i.e. variadic functions).

JavaScript allows functions to work flexibly when the arguments passed are of different types and/or at different positions:

• The useState Hook lets you pass an initial value or a function for lazy initialization, or you can skip it altogether and pass nothing to it.
• The query API from node-postgres accepts an optional callback function and returns a promise when the callback function isn’t provided.
• The write function of the file system API in Node.js defines the first argument to be either a buffer of data or a string that we write to a file.
• The extend API from the package node-extend enables deep copying by allowing you to pass an optional boolean flag as the first argument to the function call.

It’s been a running theme in my TypeScript career: I have to create and type functions like these, and I’ve definitely struggled. But, I’ve found a few solutions that have worked for me. In this post, I’ll walk you through some techniques I use to type polymorphic functions more easily.

Union type#

Union types are probably the first, and most obvious, tool you want to reach for when typing a function that accepts arguments of different types. For example:

declare function foo(a: string | boolean)  

The argument’s type could either be a string or boolean, so we use a union type to model this. Then, we use type guards inside the function body to narrow it down to its single type, i.e., string or boolean.

Now, let’s say the return value’s type depends on which specific union member the argument’s type is. How should we then go about typing it? We can represent the types of the arguments using generic types. Then, we pass them to conditional types to retrieve the right type of return value.

Let’s consider a function that generates a random integer from 0-9 when called with the string int. Or, it generates a random English letter from a - z when called with the string char.

Here’s how I would write it in JavaScript:

function getRandom(str){
  if (str === "int") {
    // generate a random integer
    return Math.floor(Math.random() * 10);
  } else {
    // generate a random char
    return String.fromCharCode(97+Math.floor(Math.random() * 26))
  }
}

To properly type this in TypeScript, follow these steps:

1. The argument str has the string union type "int" | "char" , and to make the return value’s type depend on the argument type, we must use a generic type T to represent it.
   1. function getRandom<T extends'char' | 'int'>(str: T)
2. Pass T to a generic conditional type GetReturnType to get the respective type for the return value.
   1. type GetReturnType<T> = T extends 'char' ? string : T extends 'int' ? number : never

Putting these together we have:

type GetReturnType<T> = T extends 'char' ? string : T extends 'int' ? number : never

function getRandom<T extends'char' | 'int'>(str: T): GetReturnType<T> {
  if (str === 'int') {
    // generate a random number
    return Math.floor(Math.random() * 10)  as GetReturnType<T>
  } else {
    // generate a random char
    return String.fromCharCode(97+Math.floor(Math.random() * 26))  as GetReturnType<T>
  }
}

You might be wondering about the type assertion after each return statement. I’ll explain this later.

Now, let’s say we must expand our getRandom function to also support random boolean generation.

First, we must add another union member bool to our string union type for the argument. That’s easy. But as a result, the conditional expressions inside GetReturnType quickly gets crowded:

type GetReturnType<T> = T extends 'char'
  ? string
  : T extends 'int'
  ? number
  : T extends 'bool'
  ? boolean
  : never

function getRandom<T extends 'char' | 'int' | 'bool'>(
  str: T
): GetReturnType<T> {
  if (str === 'int') {
    // generate a random number
    return Math.floor(Math.random() * 10) as GetReturnType<T>
  } else if (str === 'char') {
    // generate a random char
    return String.fromCharCode(
      97 + Math.floor(Math.random() * 26)
    ) as GetReturnType<T>
  } else {
    // generate a random boolean
    return Boolean(Math.round(Math.random())) as GetReturnType<T>
  }
}

As you can tell, this doesn’t scale well if we keep adding more types for the function to support. Luckily, we can create a record type for indexed access with type parameter T, which we defined for our argument’s type.

// interface works as well
type ReturnTypeByInputType = {
  int: number
  char: string
  bool: boolean
}

function getRandom<T extends 'char' | 'int' | 'bool'>(
  str: T
): ReturnTypeByInputType[T] {
  if (str === 'int') {
    // generate a random number
    return Math.floor(Math.random() * 10) as ReturnTypeByInputType[T]
  } else if (str === 'char') {
    // generate a random char
    return String.fromCharCode(
      97 + Math.floor(Math.random() * 26)
    ) as ReturnTypeByInputType[T]
  } else {
    // generate a random boolean
    return Boolean(Math.round(Math.random())) as ReturnTypeByInputType[T]
  }
}

Think about the DOM API documnet.querySelector—it accepts an html element tag name and returns the respective html element. It is typed in TypeScript’s source code in the exact same way.

Optional parameters#

What about functions that take a variable number of arguments? They’re extremely common, and in JavaScript you don’t really need to do anything other than just define parameters as you normally would and check them against undefined inside the function body.

In TypeScript—you probably know this already—we can model it using optional parameters marked with ?:

declare function foo(a: string, b?: boolean)

Consequently, inside the function body, b is of the union type boolean | undefined

It’s also common for such functions to return different types of values if optional parameters are actually provided or not.

Let’s say we have a function search that fetches search results asynchronously. It accepts an optional callback function as the arguments. If the callback function is supplied, it passes the search results to it. Otherwise, it returns a promise that resolves to the search results. Here’s how you might write the function in JavaScript:

function search(query, cb) {
  const res = api(query)
  if (cb) {
    res.then((data) => cb(data))
    return
  }

  return res
}

const p = search('foo') // return a promise
const v = search('foo', (data) => {}) // void

In TypeScript, we can follow these steps to type the function:

1. First, we must mark the argument cb as an optional parameter with ?
2. Then, we represent the argument cb's type with a generic type T
3. Finally, we use extends to conditionally return the right type T extends Callback ? void : Promise<Result[]>

type Callback = (results: Result[]) => void;

function search<T extends Callback | undefined = undefined>(
  query: string,
  cb?: T
): T extends Callback ? void : Promise<Result[]>{
  const res = api(query);

  if (cb) {
    res.then((data) => cb(data));
    return undefined as void & Promise<Result[]>; // assertion needed for the same reason as `getRandom` above
  }

  return res as void & Promise<Result[]>;
}

const p = search("key"); // ✅ Promise<Result[]>
const v = search("key", (data) => {} ); // ✅ void

As you might’ve noticed, there are some common themes:

1. We use extends with conditional expressions quite a lot to determine the right return’s type. And the syntax can get complicated pretty quickly.
2. There are a lot of type assertions needed.

All of these added a lot of noise to our types. So, there might be a better alternative when it comes to type complex polymorphic functions...

Function overload#

It turns out that TypeScript supports function overload, and surprisingly, it might actually be the oldest part of TypeScript. You can trace it back to TypeScript 1.1. But unlike other features added during TypeScript’s early development—enums and namespaces come to mind, which tend to get overused (especially enums) and should be replaced by other features—from my observations, function overload is actually underused and still remains useful when needed.

I think part of the reason why function overload is so underused, is because the idea of function overload just feels unnatural to many JavaScript developers. In JavaScript, we don’t have function overload—JavaScript only allows one function with a specific name within a specific scope.

However, as a dynamically typed language, JavaScript performs type checks during runtime. This means arguments are as dynamic as we need them to be and allows us to achieve the same effect as function overload—namely, having different function implementations depending on the types and number of arguments that are invoked.

A simple example of function overload#

Let’s consider a function that accepts either a number or a string, and that converts the input to the opposite type and returns it. That means, given a number, it returns the corresponding string; given a string, it returns the corresponding number. Here’s how you can write it in JavaScript:

This example is inspired by this tweet from @TkDodo

function switchIt(input) {
	if(typeof input === 'string') return Number(input)
	else return String(input)
}

And here’s how you can type this function using generics and conditional types:

function switchIt<T extends string | number>(input: T): T extends string ? number : string {
    if (typeof input === 'string') {
        return Number(input) as string & number
    } else {
        return String(input) as string & number
    }
}

const num = switchIt('1') // has type number ✅
const str = switchIt(1) // has type string ✅

Now let’s try function overload to type this. Follow these steps:

1. Write 2 separate function signatures for each version of the overloaded function

function switchIt_overloaded(input: string): number
function switchIt_overloaded(input: number): string

1. Write the overloaded function implementation.
   1. Use a union type to encompass types of each of the overloads.
   2. Within the function body, we check the types of the arguments and manually dispatch the execution to a proper code path:

   function switchIt_overloaded(input: string): number
   function switchIt_overloaded(input: number): string
   function switchIt_overloaded(input: number | string): number | string {
       if (typeof input === 'string') {
           return Number(input) 
       } else {
           return String(input) 
       } 
   }
   

With function overload, you remove:

1. Generics and the conditional types.
2. Type assertions.

And you gain benefits like:

1. Readability, since you can clearly distinguish the possible variant of the overloaded function. The types of arguments and return values are separately and explicitly written out.
2. IDE support for overloaded functions is better.

A more complex example of function overload#

Rewind to our initial search function. Following the same steps, you can re-write it using function overload:

type Callback = (results: Result[]) => void

function search_overloaded(term: string): Promise<Result[]>
function search_overloaded(term: string, cb: Callback): void
function search_overloaded(
  term: string,
  cb?: Callback
): void | Promise<Result[]> {
  const res = api(term)

  if (cb) {
    res.then((data) => cb(data))
    return
  }

  return res
}

const p = search_overloaded('key') // ✅ Promise<Result[]>
const v = search_overloaded('key', (data) => {}) // ✅ void

Again, no convoluted conditional types and generic types, and no annoying type assertions.

One last example of function overload in React#

React's useState hook is also overloaded to make it easier to use.

If you have an initial value or a function that returns a value then the state is going to be of the type of that value:

const [state] = useState(1) // number

You can also skip passing an initial value to it, and instead specify a type. Then the state ends up being of the union type:

const [state] = useState<number>() // number | undefined

If you don't even specify a type, you will get a undefined state

const [state] = useState() // undefined

This is also done via function overload.

function useState<S = undefined>(): [S | undefined, Dispatch<SetStateAction<S | undefined>>]
function useState<S>(initialState: S | (() => S)): [S , Dispatch<SetStateAction<S | undefined>>]
function useState<S>(initialState?: S | (() => S)): [S | undefined , Dispatch<SetStateAction<S | undefined>>] {
   // ...implementation  
}

Check out the source code if you are interested.

Overloaded functions are just as unsafe (even without type assertions)#

Type assertions are often considered to be code smell, and getting rid of them by leveraging function overload might seem like a big win. However, function overload is just as unsafe as type assertions.

Let’s go back to our switchIt_overloaded example and intentionally mess up its implementation to return the wrong types:

function switch_overloaded(input: string): number
function switch_overloaded(input: number): string
function switch_overloaded(input: number | string): number | string {
     if (typeof input === 'string') {
        return input // input is still a string when it should be converted to number
    } else {
        return input // input is still a number when it should be converted to string
    } 
}

const num = switch_overloaded('1') // ❌ num's type is number but it is actually a string
const str = switch_overloaded(1) // ❌ str's type is string but it is actually a number

The TypeScript compiler only checks the function body’s code against the (overloaded) function signature but it cannot tell which if else branch is supposed to handle which individual overload. As a result, we can write code that contradicts the overloaded function signature and TypeScript can’t help us.

Function overload is just an intersection of function types#

Function overload can be thought of as a syntactic sugar for intersecting function types. Think about it—when we are overloading functions like:

function switchIt(input: string): number
function switchIt(input: number): string

Then we are saying that a value (i.e. that function) of this type (the function type/signature being overloaded) can be used both as a function of the first type/signature (input: string): number and as a function of the second type/signature (input: number): string. This effectively translates into an intersection of both function types/signatures:

type F = ((input: string) => string) & ((input: number) => number) 

const switchIt_intersection: F = (input) => {
    if (typeof input === 'string') {
        return Number(input)
    } else {
        return String(input)
    }
}

const num = switchIt_intersection(1) // ✅ has the string type
const str = switchIt_intersection('1') // ✅ has the number type

And you can also write type F in the form of interface, since the interface definition automatically merges and is implicitly intersected by the compiler:

interface F {
    (input: number): string
    (input: string): number
}

Flow is another popular JavaScript type system. But, it doesn’t (fully) support the function overload syntax in TypeScript. However, it does allow you to set overloading types for functions using intersection types, exactly like what we did above with switchIt_intersection.

function switch_overloaded(input: string): number
function switch_overloaded(input: number): string
function switch_overloaded(input: number | string): number | string { // 🤔 union, not intersection
	// ...
}

My rule of thumb#

When it comes to typing polymorphic functions in TypeScript, I normally default to using generic types (constrained to a union type) along with conditional types. I only reach for function overload when I realize the function signature’s shapes for all its variants are different enough to be defined separately and explicitly.

For example:

1. The search function above is a good candidate for function overload since the return value’s type changes depend on the number of arguments that gets passed to the function (the shapes of the function signatures are very different).
2. The getRandom function is not suitable for function overload since generic types with conditional types or indexed types are already a great tool to map input types to output types. Writing it using function overload would be extremely verbose. Functions with an excessive number of overloads can be confusing to people.

The bottom line is, whether you favour function overload or generic types with conditional types, we have to be very intentional about it and tread very carefully as neither of them is completely safe.