Jsonnet の遅延評価+オブジェクトマージを Python で再現する

Jsonnet という便利なものがあります。これは JSON を生成するための DSL で、Kubernetes のマニフェストなどの生成手段として大変有用です。その便利さは、Jsonnet の備える遅延評価の機能と、+ 演算子によるオブジェクトマージに由来すると個人的に考えています。詳しくは公式チュートリアルを見ていただければわかりますが、それぞれについて簡単に説明します。

まず遅延評価の例です。

{sum: $.a + $.b} + {a: 2, b: 2}

の評価結果が以下になります。

{sum: 4, a: 2, b: 2}

$ がそれを含む最大のオブジェクトを指す({sum: $.a + $.b} の場合だとこのオブジェクト全体を $ が指す)ことを踏まえると、 {sum: $.a + $.b} の評価時点では a も b も存在しないのでエラーになりそうなところです。しかし Jsonnet は評価戦略として遅延評価を採用しているため、$ が指すものは全体の評価時に決定されます。つまり、次のように考えることが出来ます。上の式はまず以下の形になり、

{sum: $.a + $.b, a: 2, b: 2}

これが評価されるときには $.a と $.a の値が存在するため、

{sum: 4, a: 2, b: 2}

となるわけです。

次にオブジェクトマージについてです。上の例は Jsonnet のオブジェクトマージ演算子 + を Python の dict に関する | 演算子と同じ振る舞いをすると考えても解釈できますが、Jsonnet の + はさらに多機能です。

同じキーが左項と右項にある場合、右項の値が優先され、この点は Python の | と同じですが、

{a: {b:2}} + {a: {c:3}} // {a: {c: 3}}

右項のキーに + と付けると左右の値が + に項として与えられたものが結果のオブジェクトにおける値になります。

{a: {b:2}} + {a+: {c:3}} // {a: {b: 2, c: 3}}

さて、この2つの機能は便利なので Python で dict を加工する際にも使いたいと思い、再現したくなりました。まずオブジェクトマージを再現します。

from dataclasses import dataclass

@dataclass
class Plus:
    name: str
    def __hash__(self):
        return hash((self.__class__, self.name))


def merge(left, right):
    if isinstance(left, dict) and isinstance(right, dict):
        new = dict(left)
        for qualified_k, right_v in right.items():
            if isinstance(qualified_k, Plus):
                k = qualified_k.name
            else:
                k = qualified_k
            if k not in new:
                new[k] = right_v
                continue
            if isinstance(qualified_k, Plus):
                new[k] = merge(new[k], right_v)
            else:
                new[k] = right_v
        return new
    elif isinstance(left, list) and isinstance(right, list):
        return left + right
    else:
        raise TypeError()

Plus という dataclass は Jsonnet でキーに + をつけることに対応したものです。Python でオブジェクトが dict のキーとして使えるためには hashable である必要があるので __hash__ を定義しています。

使ってみます。

>>> merge({'a': [1]}, {Plus('a'): [2]})
{'a': [1, 2]}
>>> merge({'a': {'b':[1]}}, {Plus('a'): {'b':[2]}})
{'a': {'b': [2]}}
>>> merge({'a': {'b':[1]}}, {Plus('a'): {Plus('b'):[2]}})
{'a': {'b': [1, 2]}}

出来ていますね。次に遅延評価の再現です。クロージャを使います。

from dataclasses import dataclass
import typing

@dataclass
class Thunk:
    closure: typing.Callable
    def reify(self):
        return self.closure()

def reify(obj):
    if isinstance(obj, Thunk):
        return reify(obj.reify())
    elif isinstance(obj, dict):
        return {k: reify(v) for k, v in obj.items()}
    elif isinstance(obj, dict):
        return [reify(v) for v in obj]
    return obj

reify 関数は Thunk を含みうるオブジェクトの Thunk をすべて評価する関数です。使ってみます。

>>> reify(Thunk(lambda: {'a': Thunk(lambda: 1)}))
{'a': 1}

さて、これらを組み合わせるだけで遅延評価される dict のマージができるかと言うと、出来ません。なぜなら merge 関数は Thunk を一切考慮していないからです。Thunk を考慮するように書き換えます。

def merge(left, right):
    if isinstance(left, Thunk) and isinstance(right, Thunk):
        return Thunk(lambda: merge(left.reify(), right.reify()))
    elif isinstance(left, Thunk):
        return Thunk(lambda: merge(left.reify(), right))
    elif isinstance(right, Thunk):
        return Thunk(lambda: merge(left, right.reify()))
    elif isinstance(left, dict) and isinstance(right, dict):
        new = dict(left)
        for qualified_k, right_v in right.items():
            if isinstance(qualified_k, Plus):
                k = qualified_k.name
            else:
                k = qualified_k
            if k not in new:
                new[k] = right_v
                continue
            if isinstance(qualified_k, Plus):
                new[k] = merge(new[k], right_v)
            else:
                new[k] = right_v
        return new
    elif isinstance(left, list) and isinstance(right, list):
        return left + right
    else:
        raise TypeError()

ポイントは、左項もしくは右項が Thunk なら結果も Thunk になるという点です。使ってみます。

>>> reify(merge({'a': Thunk(lambda: [1])}, Thunk(lambda: {Plus('a'): Thunk(lambda: [2])})))
{'a': [1, 2]}