Python處理yaml和巢狀資料結構技巧範例

2022-06-23 14:00:55

引言

分享一些用Python處理yaml和巢狀資料結構的的一些技巧，首先從修改yaml格式檔案的問題出發，演變出了各個解決辦法，又從最後的解決辦法中引申出了普適性更強的巢狀資料結構的定位方法。

保留註釋修改yaml檔案
定位巢狀資料結構
定位巢狀資料結構2

保留註釋修改yaml檔案

yaml比之json檔案的其中一個區別就是可以註釋，這些註釋有時候是很重要的內容，就像程式碼中的註釋一樣，如果是手動編輯自然是沒有問題的，那麼如何在保留註釋的情況下用程式碼修改yaml檔案呢？

假設我們要修改的yaml檔案如下:

# 主要維護人
name: zhangsan

# 各叢集運維人員
cluster1:
  node1:
    tomcat: user11

cluster2:
  node1:
    tomcat: user21

不保留註釋

為了演示處理yaml的各個方法，這裡把不保留註釋的方法也納入到本文了。

def ignore_comment():
    data = yaml.load(text, Loader=yaml.Loader)
    data["name"] = "wangwu"
    print(yaml.dump(data))

輸出如下:

cluster1:
node1:
tomcat: user11
cluster2:
node1:
tomcat: user21
name: wangwu

很顯然，這不是我們要的結果，那麼就淘汰這個方法吧。

此方法只適用於不需要保留註釋的修改。

正規表示式

既然load, dump方法會丟棄註釋，那麼用正規表示式不就可以了麼，處理文字一定有正規表示式一席之地的。

假設還是將name: zhangsan改成name: wangwu。

def regex1():
    pattern = "name:s+w+"
    pat = re.compile(pattern=pattern)
    # 首先匹配到對應的字串
    sub_text = pat.findall(text)[0]
    # 根據這個字串找到在文字的位置
    start_index = text.index(sub_text)
    # 根據起始位置計算結束位置
    end_index = start_index + len(sub_text)
    print(start_index, end_index, text[start_index:end_index])

    # 將根據索引替換內容
    replace_text = "name: wangwu"
    new_text = text[:start_index] + replace_text + text[end_index:]
    print("="*10)
    print(new_text)

輸出如下:

8 22 name: zhangsan
==========
# 主要維護人
name: wangwu
# 各叢集運維人員
cluster1:
node1:
tomcat: user11
cluster2:
node1:
tomcat: user21

看起來不錯，好像能夠滿足需求，但是這裡有一個問題就是，假設修改是cluster2.node1.tomcat的值呢？

因為文字中有兩個tomcat的值，所以只是通過正規表示式不能一擊即中，需要多一些判斷條件，比如首先找到cluster2的起始位置，然後過濾掉小於這個起始位置的索引值，但是如果還有cluster3,cluster4呢？總的來說還是需要人工的過一遍，然後根據觀察結果來編寫正規表示式，但是這樣太不智慧，太不自動了。

此方法適用於比較容器匹配的文字。

語法樹

其實整個文字的資料結構大致如下:

無論是程式語言還是資料文字，如json, yaml, toml都可以得到這樣的語法樹，通過搜尋這顆語法樹，我們就能找到對應的鍵值對。

def tree1():
tree = yaml.compose(text)
print(tree)

輸出如下:

MappingNode(tag='tag:yaml.org,2002:map', value=[(ScalarNode(tag='tag:yaml.org,2002:str', value='name'), ScalarNode(tag='tag:yaml.org,2002:str', value='zhangsan')), (ScalarNode(tag='tag:yaml.org,2002:str', value='cluster1'), MappingNode(tag='tag:yaml.org,2002:map', value=[(ScalarNode(tag='tag:yaml.org,2002:str', value='node1'), MappingNode(tag='tag:yaml.org,2002:map', value=[(ScalarNode(tag='tag:yaml.org,2002:str', value='tomcat'), ScalarNode(tag='tag:yaml.org,2002:str', value='user11'))]))])), (ScalarNode(tag='tag:yaml.org,2002:str', value='cluster2'), MappingNode(tag='tag:yaml.org,2002:map', value=[(ScalarNode(tag='tag:yaml.org,2002:str', value='node1'), MappingNode(tag='tag:yaml.org,2002:map', value=[(ScalarNode(tag='tag:yaml.org,2002:str', value='tomcat'), ScalarNode(tag='tag:yaml.org,2002:str', value='user21'))]))]))])

通過yaml.compose方法我們就能得到一顆節點樹，並且每個節點會包括該節點的文字資訊，比如起始，終止的文字索引。

通過觀察我們能找到name: zhangsan的兩個節點, 鍵name是一個ScalarNode節點, 值zhangsan也是一個ScalarNode, 所以我們可以列印一下看看是否和正規表示式的結果一致。

def tree2():
    tree = yaml.compose(text)
    key_name_node = tree.value[0][0]
    value_name_node = tree.value[0][1]
    print(key_name_node.start_mark.pointer, value_name_node.end_mark.pointer, key_name_node.value, value_name_node.value)

輸出如下:

8 22 name zhangsan

結果與正規表示式一致，所以說明這種方法可行並且準確。

得到了修改文字的索引位置，就可以替換了，這裡就不再演示了。

此方法適合保留註釋的修改，並且定位巢狀結構較之正規表示式要簡單，並且不需要人工介入。

那麼如何定位巢狀結構呢？

定位巢狀資料結構

從上一節我們瞭解到了資料結構可以抽象成一顆語法樹, 那麼利用一些樹的搜尋演演算法就可以定位到目標文字了。

這裡展示一下不包含列表節點的搜尋演演算法。

def find_slice(tree: yaml.MappingNode, keys: List[str]) -> Tuple[Tuple[int, int], Tuple[int, int]]:
    """
    找到yaml檔案中對應鍵值對的索引, 返回一個((key起始索引, key結束索引+1), (value起始索引, value結束索引+1))的元組
    暫時只支援鍵值對的尋找.

    比如:
    >>> find_slice("name: zhangsan", ["name"])
    ((0, 4), (6, 14))
    """
    if isinstance(tree, str):
        tree = yaml.compose(tree, Loader=yaml.Loader)
    assert isinstance(tree, yaml.MappingNode), "未支援的yaml格式"
    target_key = keys[0]
    for node in tree.value:
        if target_key == node[0].value:
            key_node, value_node = node
            if len(keys) == 1:
                key_pointers = (key_node.start_mark.pointer, key_node.end_mark.pointer)
                value_pointers = (value_node.start_mark.pointer, value_node.end_mark.pointer)
                return (key_pointers, value_pointers)

            return find_slice(node[1], keys[1:])

    return ValueError("沒有找到對應的值")

演演算法核心在於遞迴。

這裡的實現並沒有處理列表節點(SequenceNode)。

假設我們要找cluster1.node1.tomcat並將其值改成changed, 程式碼如下:

def tree3():
    slices = find_slice(text, ["cluster1", "node1", "tomcat"])
    value_start_index, value_end_index = slices[1]
    replace_text = "changed"
    new_text = text[:value_start_index] + replace_text + text[value_end_index:]
    print(new_text)

輸出如下

# 主要維護人
name: zhangsan

# 各叢集運維人員
cluster1:
node1:
tomcat: changed

cluster2:
node1:
tomcat: user21

上面的演演算法只能定位key-value型別的資料結構，現在在此優化一下，讓其支援序列。

def find_slice2(tree: yaml.MappingNode, keys: List[str]) -> Tuple[Tuple[int, int], Tuple[int, int]]:
    """
    找到yaml檔案中對應鍵值對的索引, 返回一個((key起始索引, key結束索引+1), (value起始索引, value結束索引+1))的元組
    暫時只支援鍵值對的尋找.

    比如:
    >>> find_slice2("name: zhangsan", ["name"])
    ((0, 4), (6, 14))
    """
    if isinstance(tree, str):
        tree = yaml.compose(tree, Loader=yaml.Loader)
    target_key = keys[0]

    assert isinstance(tree, yaml.MappingNode) or isinstance(tree, yaml.SequenceNode), "未支援的yaml格式"

    ret_key_node = None
    ret_value_node = None
    value_pointers= (-1, -1)

    if isinstance(tree, yaml.SequenceNode):
        assert isinstance(target_key, int), "錯誤的資料格式"
        # 索引可以是負索引, 比如[1,2,3][-1]
        if len(tree.value) < abs(target_key):
            raise IndexError("索引值大於列表長度")

        node = tree.value[target_key]
        if len(keys) > 1:
            return find_slice2(tree.value[target_key], keys[1:])

        if isinstance(node, yaml.MappingNode):
            ret_key_node, ret_value_node = node.value[0]
        else:
            ret_key_node = node

    if isinstance(tree, yaml.MappingNode):
        for node in tree.value:
            if target_key == node[0].value:
                key_node, value_node = node
                if len(keys) > 1:
                    return find_slice2(node[1], keys[1:])
                ret_key_node = key_node
                ret_value_node = value_node

    if ret_key_node:
        key_pointers = (ret_key_node.start_mark.pointer, ret_key_node.end_mark.pointer)

    if ret_value_node:
        value_pointers = (ret_value_node.start_mark.pointer, ret_value_node.end_mark.pointer)

    if ret_key_node:
        return (key_pointers, value_pointers)

    return ValueError("沒有找到對應的值")

假設yaml檔案如下:

# 使用者列表
users:
  - user1: wangwu
  - user2: zhangsan

# 叢集中介軟體版本
cluster:
  - name: tomcat
    version: 9.0.63
  - name: nginx
    version: 1.21.6

def tree4():
    slices = find_slice2(text2, ["cluster", 1, "version"])
    value_start_index, value_end_index = slices[1]
    replace_text = "1.22.0"
    new_text = text2[:value_start_index] + replace_text + text2[value_end_index:]
    print(new_text)

輸出如下:

# 使用者列表
users:
- user1: wangwu
- user2: zhangsan

# 叢集中介軟體版本
cluster:
- name: tomcat
version: 9.0.63
- name: nginx
version: 1.22.0

結果符合預期。

定位巢狀資料結構2

上面介紹瞭如何定位巢狀的資料結構樹，這一節介紹一下如何定位較深的樹結構(主要指python字典)。

鏈式呼叫get

在獲取api資料的時候因為想要的資料結構比較深，用索引會報錯，那麼就需要捕獲異常，這樣很麻煩，並且程式碼很冗長，比如:

data1 = {"message": "success", "data": {"limit": 0, "offset": 10, "total": 100, "data": ["value1", "value1"]}}
data2 = {"message": "success", "data": None}
data3 = {"message": "success", "data": {"limit": 0, "offset": 10, "total": 100, "data": None}}

上面的資料結構很有可能來自同一個api結構，但是資料結構卻不太一樣。

如果直接用索引，就需要捕獲異常，這樣看起來很煩，那麼可以利用字典的get方法。

ret = data1.get("data", {}).get("data", [])
if ret:
    pass # 做一些操作
if data2.get("data"):
    ret = data2["data"].get("data", [])
ret = data3.get("data", {}).get("data", [])

通過給定一個預期的資料空物件，讓get可以一致寫下去。

寫一個遞迴的get

起始在之前的find_slice方法中，我們就發現遞迴可以比較好的處理這種巢狀的資料結構，我們可以寫一個遞迴處理常式，用來處理很深的資料結構。

假設資料結構如下:

data = {"message": "success", "data": {"data": {"name": "zhangsan", "scores": {"math": {"mid-term": 88, "end-of-term": 90}}}}}

我們的目標就是獲取資料中張三期中數學成績: 88

實現的遞迴呼叫如下:

def super_get(data: Union[dict, list], keys: List[Union[str, int]]):
    assert isinstance(data, dict) or isinstance(data, list), "只支援字典和列表型別"
    key = keys[0]

    if isinstance(data, list) and isinstance(key, int):
        try:
            new_data = data[key]
        except IndexError as exc:
            raise IndexError("索引值大於列表長度") from exc
    elif isinstance(data, dict) and isinstance(key, str):
        new_data = data.get(key)
    else:
        raise ValueError(f"資料型別({type(data)})與索引值型別(f{type(key)}不匹配")

    if len(keys) == 1:
        return new_data

    if not isinstance(new_data, dict) and not isinstance(new_data, list):
        raise ValueError("找不到對應的值")
    return super_get(new_data, keys[1:])

然後執行程式碼:

def get2():
    data = {"message": "success", "data": {"data": {"zhangsan": {"scores": {"math": {"mid-term": 88, "end-of-term": 90}}}}}}
    print(super_get(data, ["data", "data", "zhangsan", "scores", "math", "mid-term"]))
    # 輸出 88
    data = {"message": "success", "data": {"data": {"zhangsan": {"scores": {"math": [88, 90]}}}}}
    print(super_get(data, ["data", "data", "zhangsan", "scores", "math", 0]))
    # 輸出 88
    data = {"message": "success", "data": {"data": {"zhangsan": {"scores": {"math": [88, 90]}}}}}
    print(super_get(data, ["data", "data", "zhangsan", "scores", "math", -1]))
    # 輸出 90