Extending Kolo with plugins

Kolo supports defining plugins that add custom data into Kolo traces.

A plugin is a Python package that has one or more entry points in pyproject.toml:

[project.entry-points."kolo.processors"]
myprocessor = "mypackage.processors:processor_config"
another = "mypackage.processors:another_config"

Each of these entry points should be a Python dictionary containing configuration data for the processor:

# mypackage/processors.py

processor_config = {
    "co_names": ["send"],
    "path_fragment": "/mypackage/client.py",
    "call": call,
    "call_type": "outbound_http_request",
    "return_type": "outbound_http_response",
    "subtype": "mypackage",
    "process": process,
    "build_context": build_context,
}

The entries are split into three categories: build_context, those used for selecting which code to trace and those used when processing that code.

build_context

This optional function allows creating a dictionary that will be passed to call and process to allow sharing data between calls. Kolo doesn’t use this dictionary for any other purpose.

type Context = Dict[str, Any]


def build_context(config: Dict[str, Any]) -> Context {
    return {"mydata": []}
}

Selecting code to trace

co_names

A list of names of code objects to match against. These are usually function or method names.

path_fragment

A fragment of a file path. This is used to match against a code object’s co_filename. Path separators should be written using / and will be converted to \ on Windows.

call

An optional function that allows for further control over when the processor runs. For example, to only run on return events you could write:

def call(
    frame: types.FrameType,
    event: str,
    arg: object,
    context: Context,
) -> bool:
    return event == "return"

frame, event and arg are described in the documentation for sys.setprofile. context is the dictionary created by build_context.

The call function is used in addition to co_names and path_fragment, not instead of them.

Processing code

call_type

This string is included in the returned data under the type key for call events. It is used by other tools (e.g. VSCode) to identify how to handle the data.

return_type

This string is included in the returned data under the type key for return events. It is used by other tools (e.g. VSCode) to identify how to handle the data.

subtype

This optional string is included in the returned data under the subtype key. It can be used to disambiguate between different tools sharing the same type.

process

An optional function that allows for further control over the returned data. For example, to include all local variables you could write:

def process(
    frame: types.FrameType,
    event: str,
    arg: object,
    context: Context,
) -> Dict[str, Any]:
    return {"locals": frame.f_locals}

frame, event and arg are described in the documentation for sys.setprofile. context is the dictionary created by build_context.

The output of process is merged into the default return data, so can override keys such as type and subtype if necessary.

Integrating with VSCode

Currently we don’t provide a way to write a plugin for the VSCode extension. Instead, plugin authors are encouraged to reuse existing features.

Background jobs

To implement a processor for a background job runner like celery or huey:

• Set call_type to "background_job".

• Set return_type to "background_job_end".

• Set subtype to the name of your job runner.

• Add a process handler that returns a dictionary with name, args and kwargs keys. name is the name of the background job and args and kwargs are the data passed to it.

Outbound http requests

To implement a processor for a http library like requests, httpx, urllib3 or urllib:

• Set call_type to "outbound_http_request".

• Set return_type to "outbound_http_response".

• Set subtype to the name of the http library.

• Add a process handler:

def process(
    frame: types.FrameType,
    event: str,
    arg: Any,
    context: Context,
) -> Dict[str, Any]:
    if event == "call":
        return {
            "body": request_body,
            "headers": request_headers,
            "method": method,
            "method_and_full_url": f"{method} {url}",
            "url": url
        }
    elif event == "return":
        return {
            "body": response_body,
            "headers": response_headers,
            "method": method,
            "method_and_full_url": f"{method} {url}",
            "status_code": status_code,
            "url": url
        }

Local processor plugins

You can also define a custom processor locally. The process is identical, except you edit .kolo/config.toml instead:

processors = [
    "mypackage.processors:processor_config",
    "mypackage.processors:another_config",
]