Anyone using cdk-nag in python and creating custom rules?

Documentation for it is all in js/ts and somethings just straight up do not work the same. In ts its expecting a rule function callback, for example, but passing a function reference in python results in an error.

Link to docs

This is my initial stab at making a custom rule

Current error:

RuntimeError: @jsii/kernel.SerializationError: Passed to parameter params of method cdk-nag.NagPack#applyRule: Unable to deserialize value as cdk-nag.IApplyRule

Code:

from aws_cdk import aws_kms as kms
from cdk_nag import IApplyRule, NagMessageLevel, NagPack, NagRuleCompliance
from constructs import IConstruct
from typing import cast


def kms_must_have_alias_rule_fn(node: IConstruct) -> NagRuleCompliance:
    if isinstance(node, kms.Key):
        if node.alias is None:
            return NagRuleCompliance.NON_COMPLIANT
        else:
            return NagRuleCompliance.COMPLIANT
    else:
        return NagRuleCompliance.NOT_APPLICABLE


class CisChecks(NagPack):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._pack_name = 'CisChecks'

    def visit(self, node: IConstruct) -> None:
        self._apply_rule(cast(IApplyRule,
                              {
                                  'ruleSuffixOverride': 'KMSKeyMustHaveAlias',
                                  'info': 'CIS requires that all KMS keys have an alias',
                                  'explanation': 'It is really helpful when managing keys',
                                  'level': NagMessageLevel.ERROR,
                                  'rule': kms_must_have_alias_rule_fn(node),
                                  'node': node,
                              }
                              )
                         )

Since none of the docs are in python, I have no idea what the shape of the IApplyRule interface is supposed to be.