Sử dụng thẻ tùy chỉnh trong PyYAML để lấy giá trị biến môi trường

Giới thiệu

YAML thường được sử dụng cho các tệp cấu hình. Khi cần cấu hình tên người dùng và mật khẩu trong tệp YAML, việc ghi trực tiếp các thông tin này và tải lên kho lưu trữ mã nguồn có thể dễ dàng dẫn đến rò rỉ mật khẩu.

Để giải quyết vấn đề này, chúng ta có thể sử dụng thẻ tùy chỉnh (custom tag) trong PyYAML. Trong PyYAML, mỗi thẻ đại diện cho một kiểu dữ liệu, và chúng ta có thể định nghĩa thẻ riêng của mình để đọc biến môi trường.

Định nghĩa thẻ tùy chỉnh

Chúng ta sẽ định nghĩa một thẻ mới là !env và viết một hàm xử lý tương ứng (trong PyYAML được gọi là constructor). Tệp cấu hình YAML:

user: !env ${USER}  # Biểu thị biến môi trường USER, tức là tên người dùng hiện tại

Tệp Python:

import os
import yaml  # Cần cài đặt pip install pyyaml

def bien_moi_truong_constructor(loader, node):
    gia_tri = loader.construct_scalar(node)   # Phương pháp cố định của PyYAML loader, dùng để xây dựng giá trị biến từ nút hiện tại
    ten_bien = gia_tri.strip('${} ')  # Loại bỏ các ký tự đặc biệt và khoảng trắng ở đầu và cuối giá trị biến (ví dụ ${USER})
    return os.getenv(ten_bien, gia_tri)  # Thử lấy giá trị tương ứng với tên biến trong môi trường, nếu không lấy được thì sử dụng giá trị mặc định (ban đầu là ${USER})

yaml.SafeLoader.add_constructor('!env', bien_moi_truong_constructor)  # Thêm thẻ mới và constructor cho SafeLoader

with open('demo.yml') as f:  
    print(yaml.safe_load(f))  # Mở tệp và sử dụng SafeLoader để tải nội dung

Kết quả:

{'user': 'ten_nguoi_dung'}

Phân phối mẫu khớp cho thẻ

Hiện tại, trong tệp YAML, biến môi trường chỉ có thể được sử dụng với kiểu khai báo bắt buộc !env ${tên_biến}. Nếu muốn sử dụng trực tiếp ${tên_biến}, chúng ta cần chỉ định một mẫu biểu thức chính quy cho thẻ này, tức là tự động sử dụng thẻ !env khi nhận định dạng tương tự ${tên_biến}.

Tệp cấu hình YAML:

user: !env ${USER}  # Biểu thị biến môi trường USER, tức là tên người dùng hiện tại
path: ${PATH}  # Mong muốn có thể sử dụng trực tiếp

Tệp Python:

import os
import re
import yaml

pattern = re.compile('\${\w+}')  # Khớp với ${một hoặc nhiều chữ cái hoặc số}

def bien_moi_truong_constructor(loader, node):
    gia_tri = loader.construct_scalar(node)  
    ten_bien = gia_tri.strip('${} ')  
    return os.getenv(ten_bien, gia_tri)  

yaml.SafeLoader.add_constructor('!env', bien_moi_truong_constructor)  # Thêm thẻ mới và constructor tương ứng
yaml.SafeLoader.add_implicit_resolver('!env', pattern, None)    # Chỉ định mẫu biểu thức chính quy cho thẻ

with open('demo.yml') as f:  
    print(yaml.safe_load(f))  # Mở tệp và sử dụng SafeLoader để tải nội dung

Kết quả:

{'user': 'ten_nguoi_dung', 'path': '...<bỏ qua>'}

Sử dụng nhiều biến trong một nút

Nếu muốn sử dụng nhiều biến trong một nút, ví dụ:

user: !env ${USER}
path: ${PATH}
msg: Tên người dùng hiện tại ${USER} đường dẫn hệ thống ${PATH}

Chúng ta cần thay thế từng giá trị trong giá trị nút (dạng chuỗi). Đầu tiên, chúng ta cần sửa đổi mẫu khớp để cho phép ${biến} có thể có nhiều ký tự tùy ý xung quanh:

pattern = re.compile('.*?(\${\w+}).*?')  # Xung quanh có thể có nhiều ký tự tùy ý, sử dụng dấu ngoặc đơn để chỉ lấy nội dung biến ${tên_bien}, `?` biểu thị khớp không tham lam.

Mã đầy đủ:

import os
import re
import yaml

pattern = re.compile('.*?(\${\w+}).*?')  # Xung quanh có thể có nhiều ký tự tùy ý, sử dụng dấu ngoặc đơn để chỉ lấy nội dung biến ${tên_bien}, `?` biểu thị khớp không tham lam.

def bien_moi_truong_constructor(loader, node):
    gia_tri = loader.construct_scalar(node)
    for item in pattern.findall(gia_tri):  # Lặp qua tất cả các biến ${tên_bien} được khớp
        ten_bien = item.strip('${} ')    # Ví dụ: USER
        gia_tri = gia_tri.replace(item, os.getenv(ten_bien, item))  # Thay thế biến bằng giá trị tương ứng từ môi trường
    return gia_tri                                  # Ví dụ: thay thế ${USER} bằng superin, nếu không lấy được thì sử dụng giá trị gốc ${USER}

yaml.SafeLoader.add_constructor('!env', bien_moi_truong_constructor)
yaml.SafeLoader.add_implicit_resolver('!env', pattern, None)

with open('demo.yml') as f:
    print(yaml.safe_load(f))

Kết quả:

{'user': 'ten_nguoi_dung', 'path': '...<bỏ qua>', 'msg': 'Tên người dùng hiện tại ten_nguoi_dung đường dẫn hệ thống ...<bỏ qua>'}

Thẻ: PyYAML biến môi trường yaml Cấu hình thẻ tùy chỉnh

Đăng vào ngày 28 tháng 6 lúc 07:03