加载 YAML 文件后添加类型提示的正确方法是什么？

Question

我正在向我的 Python 代码添加类型提示，并且想知道对加载的 YAML 文件进行类型提示的正确方法是什么，因为它是任意数量字典的字典。

有没有比类型提示更好的方法返回加载的 YAML 文件Dict[str, Dict[str, Any]]？

这是函数：

def load_yaml(yaml_in: str) -> Dict[str, Dict[str, Any]]:
    return yaml.load(open(yaml_in), Loader=yaml.FullLoader)

以下是正在加载的 YAML 文件的示例：

VariableMap:
    var1: 'time'
    var2: 'param_name'

GlobalVariables:
    limits:
        x-min:
        x-max:
        y-min:
        y-max:

Plots:
    plot1:
        file: 
        x_data: 'date'
        y_data: [{param: 'param1', label: "param1", color: 'red', linestyle: '-'},
                 {param: 'param2', label: "param2", color: 'black', linestyle: '--'}]
        labels:
            title: {label: 'title', fontsize: '9'}
            x-axis: {xlabel: 'x-label', fontsize: '9'}
            y-axis: {ylabel: 'y-label', fontsize: '9'}
        limits:
            x-min: 0
            x-max: 100
            y-min:
            y-max:

Figures:
    fig1:
        shape: [1, 1]
        size: [6, 8]
        plots: ['plot1']

Answer 1

我将为此使用 dataclass 和dataclass_json 。那么返回类型可以是一个实际的类。您还可以使用该LetterCase选项将 yaml 文件中的格式映射到默认的 Python 命名方案 - 蛇形命名法。

像这样的东西：

from dataclasses import dataclass
from dataclasses_json import dataclass_json, LetterCase, Undefined
from typing import Optional

@dataclass_json
@dataclass
class VariableMap:
    var1: str
    var2: str


@dataclass_json(letter_case=LetterCase.KEBAB, )
@dataclass
class Limits:
    x_min: Optional[int]
    x_max: Optional[int]
    y_min: Optional[int]
    y_max: Optional[int]

@dataclass_json
@dataclass
class GlobalVariables:
    limits: Limits

@dataclass_json
@dataclass
class TitleLabel:
    label: str
    fontsize: str

@dataclass_json
@dataclass
class XAxisLabel:
    xlabel: str
    fontsize: str

@dataclass_json
@dataclass
class YAxisLabel:
    ylabel: str
    fontsize: str

@dataclass_json(letter_case=LetterCase.KEBAB)
@dataclass
class Labels:
    title: TitleLabel
    x_axis: XAxisLabel
    y_axis: YAxisLabel


@dataclass_json
@dataclass
class Param:
    param: str
    label: str
    color: str
    linestyle: str

@dataclass_json(undefined=Undefined.EXCLUDE)
@dataclass
class Plot:
    file: Optional[str]
    x_data: str
    y_data: list[Param]
    labels: Labels
    limits: Limits

@dataclass_json
@dataclass
class Figure:
    shape: list[int]
    size: list[int]
    plots: list[str]

@dataclass_json
@dataclass
class Plots:
    plot1: Plot

@dataclass_json
@dataclass
class Figures:
    fig1: Figure

@dataclass_json(letter_case=LetterCase.PASCAL, undefined=Undefined.EXCLUDE)
@dataclass
class YamlRoot:
    variable_map: VariableMap
    global_variables: GlobalVariables
    plots: Plots
    figures: Figures


import yaml

def load_yaml(yaml_content: str) -> YamlRoot:
    d = yaml.safe_load(yaml_content)
    return YamlRoot.from_dict(d, infer_missing=True)

if __name__ == "__main__":
    
    yaml_text = """
VariableMap:
    var1: 'time'
    var2: 'param_name'

GlobalVariables:
    ...
"""
    
    print(load_yaml(yaml_text))


"""
Result (cleaned for readability):

YamlRoot(
  variable_map=VariableMap(var1='time', var2='param_name'),
  global_variables=GlobalVariables(
    limits=Limits(x_min=None, x_max=None, y_min=None, y_max=None)
  ),
  plots=Plots(
    plot1=Plot(
      file=None,
      x_data='date',
      y_data=[
        Param(param='param1', label='param1', color='red', linestyle='-'),
        Param(param='param2', label='param2', color='black', linestyle='--')
      ],
      labels=Labels(
        title=TitleLabel(label='title', fontsize='9'),
        x_axis=XAxisLabel(xlabel='x-label', fontsize='9'),
        y_axis=YAxisLabel(ylabel='y-label', fontsize='9')
      ),
      limits=Limits(x_min=0, x_max=100, y_min=None, y_max=None)
    )
  ),
  figures=Figures(
    fig1=Figure(shape=[1, 1], size=[6, 8], plots=['plot1'])
  )
)
"""