Chalk home page
  1. Quickstart
  2. Editor Plugin

Mypy plugin

Chalk now recommends using chalk stubgen tool over the MyPy plugin.

Chalk extensively uses the Python type system to allow you to naturally express your feature pipelines. Chalk's type annotations and feature set constructors are not natively supported by editors, which can trigger warnings and errors in your development environment.

Fortunately, Chalk ships with a mypy plugin for providing advanced support and type checking for your feature pipelines. We support mypy via its plugin system.


At the root of your project, create a file called mypy.ini with the following contents:

plugins = chalk.mypy_plugin

(if you created your project with the Chalk CLI, this will have been done for you already.)

You also need to install a mypy plugin for your editor. For example, for PyCharm and IntelliJ, we recommend the Mypy Plugin.

Plugin capabilities

Incomplete feature sets

Return only part of a feature set from a resolver.

class UserFeatures:
    name: str
    bday: str
    age: int

def incomplete_set() -> Features[, UserFeatures.age]:
    # We're omitting `age` and returning only `name` and `bday`
    return UserFeatures(name="joseph", age=34)

Chalk allows you to return only a subset of the fields for a particular feature set. In the above example, UserFeatures contains three fields, but our resolver returns only two of them. Our plugin converts

  • UserFeatures(name=..., bday=...) to
  • Features[, UserFeatures.bday]

Unordered output

Specified the returned features in any order.

def unordered_output() -> Features[UserFeatures.age,]:
    # Order of arguments need not line up. `Features` is treated as a set.
    return UserFeatures(name="joseph", age=34)

In this example, we have flipped the order of the arguments in the type signature from the order of the arguments given to UserFeatures. Because the plugin sorts the types before comparing, this example is also valid.

Type-mismatch on features

Recognize when you're returning the wrong features.

def name_age_bad() -> Features[, UserFeatures.age]:
    # We're getting Features[name, bday] and expected Features[name, age]
    return UserFeatures(name="4", bday="")

Chalk can check that the subset of features returned from a function match the intended returned features. This example results in the following error message:

error: Incompatible return value type:
         Received: Features[bday, name]
         Expected: Features[age,  name]

Here, we’re returning the and UserFeatures.bday features, but we expected to have the UserFeatures.age feature instead of the UserFeatures.bday feature.

Feature value type checking

Check that feature values receive correct types.

def name_bad_type() -> Features[]:
    # Knows that name should be a string, not an int
    return UserFeatures(name=4)

Mypy surfaces the following error for this example:

error: Argument "name" to "UserFeatures" has incompatible type "int"; expected "str"

Here, should be given as a string, but we have instead supplied an int.

Constructor arguments

Check arguments to constructor of feature sets.


Mypy surfaces the following error for this example:

error: Unexpected keyword argument "bogus_field" for "UserFeatures"

UserFeatures.bogus_field is not a feature on the UserFeatures class, and so cannot be used in the __init__ method.