532 lines
19 KiB
Python
532 lines
19 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""Ti-1 氯化车间质量预测 · 模型训练与评估（Issue #69 / PRD §5.3 ①）。
|
|||
|
|
|
|||
|
|
承接 #68（特征工程）产出的 ``FeatureMatrix``：把"特征矩阵 + 质量标签 → 可评估、
|
|||
|
|
可解释的预测模型"这条链路**模板化、可测试**，且与 #73 模型部署解耦。
|
|||
|
|
|
|||
|
|
PRD 设计口径
|
|||
|
|
------------
|
|||
|
|
- 架构表（PRD §5.3）：``质量预测 | 预测 | 入：特征矩阵；
|
|||
|
|
出：质量指标预测值(纯度/杂质) | ① 质量预测 | 中``。
|
|||
|
|
- 模板化技术路径：模型超参（``alpha`` 正则强度、``target`` 标签列）外置为
|
|||
|
|
JSON/YAML 超参包，切换模板只改超参包（PRD §5.3「换行业只改 Recipe」）。
|
|||
|
|
- 验收（PRD §10 DoD）：质量预测在客户数据上达约定 R² / MAE 指标。
|
|||
|
|
|
|||
|
|
本模块交付
|
|||
|
|
----------
|
|||
|
|
1. **岭回归 ``RidgeRegression``**：纯标准库最小二乘 + L2 正则（闭式解），
|
|||
|
|
``fit(X, y)`` / ``predict(X)`` / ``coef_`` / ``intercept_``。不依赖 numpy。
|
|||
|
|
2. **质量预测模型 ``QualityModel``**：聚合特征名 + 目标列 + 岭回归，提供
|
|||
|
|
``fit(matrix, target)`` / ``predict(matrix)`` / ``evaluate(matrix, target)``
|
|||
|
|
（R² / MAE / RMSE）/ ``explain()``（权重 → 特征贡献，可溯源）。
|
|||
|
|
3. **超参包 ``ModelRecipe``**：``target``/``alpha``/``feature_names`` 外置
|
|||
|
|
JSON/YAML 加载（零依赖 YAML 子集解析）。
|
|||
|
|
4. **评估指标**：R²、MAE、RMSE 纯标准库实现。
|
|||
|
|
|
|||
|
|
设计要点
|
|||
|
|
--------
|
|||
|
|
- **零运行时依赖**（纯标准库）：矩阵运算手写（不依赖 numpy/sklearn）。
|
|||
|
|
- **可解释**：``explain()`` 输出每个特征的权重 × 方差贡献度，供 #73 接入驾驶舱
|
|||
|
|
展示"为何预测这个纯度"（对齐 PRD"要求结果可解释、可溯源"）。
|
|||
|
|
- **稳健**：L2 正则避免共线性/过拟合；数据不足时返回明确错误而非崩溃。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import math
|
|||
|
|
import os
|
|||
|
|
from dataclasses import dataclass, field
|
|||
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
|||
|
|
|
|||
|
|
NAN = float("nan")
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelError(ValueError):
|
|||
|
|
"""质量预测模型错误（数据不足 / 维度不匹配 / 奇异 等）。"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 纯标准库线性代数（最小二乘岭回归闭式解）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _matmul_at_a(a: Sequence[Sequence[float]]) -> List[List[float]]:
|
|||
|
|
"""计算 AᵀA（n×n）。"""
|
|||
|
|
n = len(a[0]) if a else 0
|
|||
|
|
res = [[0.0] * n for _ in range(n)]
|
|||
|
|
for row in a:
|
|||
|
|
for i in range(n):
|
|||
|
|
ri = row[i]
|
|||
|
|
if ri == 0.0:
|
|||
|
|
continue
|
|||
|
|
for j in range(n):
|
|||
|
|
res[i][j] += ri * row[j]
|
|||
|
|
return res
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _matvec_at_b(a: Sequence[Sequence[float]], b: Sequence[float]) -> List[float]:
|
|||
|
|
"""计算 Aᵀb（n）。"""
|
|||
|
|
n = len(a[0]) if a else 0
|
|||
|
|
res = [0.0] * n
|
|||
|
|
for row, y in zip(a, b):
|
|||
|
|
for i in range(n):
|
|||
|
|
res[i] += row[i] * y
|
|||
|
|
return res
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _solve(A: List[List[float]], b: List[float]) -> List[float]:
|
|||
|
|
"""高斯消元解 Ax=b（带部分主元）。A 会被修改。"""
|
|||
|
|
n = len(A)
|
|||
|
|
# 增广
|
|||
|
|
M = [list(A[i]) + [b[i]] for i in range(n)]
|
|||
|
|
for col in range(n):
|
|||
|
|
# 主元
|
|||
|
|
pivot = max(range(col, n), key=lambda r: abs(M[r][col]))
|
|||
|
|
if abs(M[pivot][col]) < 1e-12:
|
|||
|
|
raise ModelError("矩阵奇异（特征共线或数据不足），无法求解")
|
|||
|
|
M[col], M[pivot] = M[pivot], M[col]
|
|||
|
|
piv = M[col][col]
|
|||
|
|
for j in range(col, n + 1):
|
|||
|
|
M[col][j] /= piv
|
|||
|
|
for r in range(n):
|
|||
|
|
if r == col:
|
|||
|
|
continue
|
|||
|
|
factor = M[r][col]
|
|||
|
|
if factor == 0.0:
|
|||
|
|
continue
|
|||
|
|
for j in range(col, n + 1):
|
|||
|
|
M[r][j] -= factor * M[col][j]
|
|||
|
|
return [M[i][n] for i in range(n)]
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 岭回归
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
class RidgeRegression:
|
|||
|
|
"""岭回归（L2 正则最小二乘，闭式解）。纯标准库。
|
|||
|
|
|
|||
|
|
解：``w = (XᵀX + αI)⁻¹ Xᵀy``。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, alpha: float = 1.0):
|
|||
|
|
if alpha < 0:
|
|||
|
|
raise ModelError(f"alpha 不能为负: {alpha}")
|
|||
|
|
self.alpha = alpha
|
|||
|
|
self.coef_: List[float] = []
|
|||
|
|
self.intercept_: float = 0.0
|
|||
|
|
self._n_features: int = 0
|
|||
|
|
|
|||
|
|
def fit(self, X: Sequence[Sequence[float]], y: Sequence[float]) -> "RidgeRegression":
|
|||
|
|
m = len(X)
|
|||
|
|
if m == 0:
|
|||
|
|
raise ModelError("训练集为空")
|
|||
|
|
n = len(X[0])
|
|||
|
|
if n == 0:
|
|||
|
|
raise ModelError("特征数为 0")
|
|||
|
|
if len(y) != m:
|
|||
|
|
raise ModelError(f"X/y 行数不匹配: {m} != {len(y)}")
|
|||
|
|
self._n_features = n
|
|||
|
|
# 中心化（数值稳定 + 让 intercept 可独立）
|
|||
|
|
x_mean = [sum(X[i][j] for i in range(m)) / m for j in range(n)]
|
|||
|
|
y_mean = sum(y) / m
|
|||
|
|
Xc = [[X[i][j] - x_mean[j] for j in range(n)] for i in range(m)]
|
|||
|
|
yc = [y[i] - y_mean for i in range(m)]
|
|||
|
|
# 正规方程 (XᵀX + αI) w = Xᵀy
|
|||
|
|
A = _matmul_at_a(Xc)
|
|||
|
|
for i in range(n):
|
|||
|
|
A[i][i] += self.alpha
|
|||
|
|
b = _matvec_at_b(Xc, yc)
|
|||
|
|
self.coef_ = _solve(A, b)
|
|||
|
|
self.intercept_ = y_mean - sum(self.coef_[j] * x_mean[j] for j in range(n))
|
|||
|
|
return self
|
|||
|
|
|
|||
|
|
def predict(self, X: Sequence[Sequence[float]]) -> List[float]:
|
|||
|
|
if not self.coef_:
|
|||
|
|
raise ModelError("模型未训练")
|
|||
|
|
return [self.intercept_ + sum(self.coef_[j] * row[j]
|
|||
|
|
for j in range(self._n_features))
|
|||
|
|
for row in X]
|
|||
|
|
|
|||
|
|
def to_dict(self) -> Dict[str, Any]:
|
|||
|
|
return {"alpha": self.alpha, "coef": list(self.coef_),
|
|||
|
|
"intercept": self.intercept_,
|
|||
|
|
"n_features": self._n_features}
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_dict(cls, d: Dict[str, Any]) -> "RidgeRegression":
|
|||
|
|
m = cls(alpha=float(d.get("alpha", 1.0)))
|
|||
|
|
m.coef_ = [float(c) for c in d.get("coef", [])]
|
|||
|
|
m.intercept_ = float(d.get("intercept", 0.0))
|
|||
|
|
m._n_features = int(d.get("n_features", len(m.coef_)))
|
|||
|
|
return m
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 评估指标
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def r2_score(y_true: Sequence[float], y_pred: Sequence[float]) -> float:
|
|||
|
|
"""决定系数 R²。"""
|
|||
|
|
if len(y_true) != len(y_pred):
|
|||
|
|
raise ModelError("y_true/y_pred 长度不匹配")
|
|||
|
|
n = len(y_true)
|
|||
|
|
if n == 0:
|
|||
|
|
return NAN
|
|||
|
|
mean = sum(y_true) / n
|
|||
|
|
ss_res = sum((y_true[i] - y_pred[i]) ** 2 for i in range(n))
|
|||
|
|
ss_tot = sum((y_true[i] - mean) ** 2 for i in range(n))
|
|||
|
|
if ss_tot == 0:
|
|||
|
|
return 1.0 if ss_res == 0 else 0.0
|
|||
|
|
return 1.0 - ss_res / ss_tot
|
|||
|
|
|
|||
|
|
|
|||
|
|
def mae_score(y_true: Sequence[float], y_pred: Sequence[float]) -> float:
|
|||
|
|
if len(y_true) != len(y_pred):
|
|||
|
|
raise ModelError("y_true/y_pred 长度不匹配")
|
|||
|
|
n = len(y_true)
|
|||
|
|
return sum(abs(y_true[i] - y_pred[i]) for i in range(n)) / n if n else NAN
|
|||
|
|
|
|||
|
|
|
|||
|
|
def rmse_score(y_true: Sequence[float], y_pred: Sequence[float]) -> float:
|
|||
|
|
if len(y_true) != len(y_pred):
|
|||
|
|
raise ModelError("y_true/y_pred 长度不匹配")
|
|||
|
|
n = len(y_true)
|
|||
|
|
if n == 0:
|
|||
|
|
return NAN
|
|||
|
|
return math.sqrt(sum((y_true[i] - y_pred[i]) ** 2 for i in range(n)) / n)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 质量预测模型（聚合特征 + 目标 + 岭回归）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class ModelRecipe:
|
|||
|
|
"""模型超参包（外置 JSON/YAML）。"""
|
|||
|
|
|
|||
|
|
target: str # 质量标签列名（如 TiCl₄纯度）
|
|||
|
|
alpha: float = 1.0 # L2 正则强度
|
|||
|
|
feature_names: List[str] = field(default_factory=list)
|
|||
|
|
unit: str = ""
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_dict(cls, d: Dict[str, Any]) -> "ModelRecipe":
|
|||
|
|
return cls(
|
|||
|
|
target=str(d.get("target", "")).strip(),
|
|||
|
|
alpha=float(d.get("alpha", 1.0)),
|
|||
|
|
feature_names=[str(x) for x in d.get("feature_names", [])],
|
|||
|
|
unit=str(d.get("unit", "")).strip(),
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def validate(self) -> List[str]:
|
|||
|
|
errs = []
|
|||
|
|
if not self.target:
|
|||
|
|
errs.append("target 不能为空")
|
|||
|
|
if self.alpha < 0:
|
|||
|
|
errs.append(f"alpha 不能为负: {self.alpha}")
|
|||
|
|
return errs
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class Evaluation:
|
|||
|
|
"""评估结果。"""
|
|||
|
|
|
|||
|
|
r2: float
|
|||
|
|
mae: float
|
|||
|
|
rmse: float
|
|||
|
|
n_samples: int
|
|||
|
|
|
|||
|
|
def to_dict(self) -> Dict[str, Any]:
|
|||
|
|
return {"r2": round(self.r2, 6), "mae": round(self.mae, 6),
|
|||
|
|
"rmse": round(self.rmse, 6), "n_samples": self.n_samples}
|
|||
|
|
|
|||
|
|
def passes(self, *, min_r2: float = 0.0, max_mae: float = math.inf) -> bool:
|
|||
|
|
return (self.r2 >= min_r2 and self.mae <= max_mae
|
|||
|
|
and not math.isnan(self.r2))
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QualityModel:
|
|||
|
|
"""质量预测模型：特征名 + 目标列 + 岭回归。"""
|
|||
|
|
|
|||
|
|
def __init__(self, recipe: ModelRecipe):
|
|||
|
|
errs = recipe.validate()
|
|||
|
|
if errs:
|
|||
|
|
raise ModelError("超参包校验失败: " + "; ".join(errs))
|
|||
|
|
self.recipe = recipe
|
|||
|
|
self.regression: Optional[RidgeRegression] = None
|
|||
|
|
self._feature_std: List[float] = []
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def fitted(self) -> bool:
|
|||
|
|
return self.regression is not None
|
|||
|
|
|
|||
|
|
def fit(self, X: Sequence[Sequence[float]], y: Sequence[float]) -> "QualityModel":
|
|||
|
|
"""X 行=样本，列=特征（顺序与 recipe.feature_names 对齐）。"""
|
|||
|
|
self.regression = RidgeRegression(alpha=self.recipe.alpha).fit(X, y)
|
|||
|
|
# 记录训练集每列标准差，供 explain() 计算尺度归一化重要性
|
|||
|
|
m = len(X)
|
|||
|
|
n = len(X[0]) if X else 0
|
|||
|
|
if m > 1 and n:
|
|||
|
|
means = [sum(X[i][j] for i in range(m)) / m for j in range(n)]
|
|||
|
|
self._feature_std = [
|
|||
|
|
math.sqrt(sum((X[i][j] - means[j]) ** 2 for i in range(m)) / (m - 1))
|
|||
|
|
for j in range(n)]
|
|||
|
|
else:
|
|||
|
|
self._feature_std = [1.0] * n
|
|||
|
|
return self
|
|||
|
|
|
|||
|
|
def predict(self, X: Sequence[Sequence[float]]) -> List[float]:
|
|||
|
|
if self.regression is None:
|
|||
|
|
raise ModelError("模型未训练，先 fit()")
|
|||
|
|
return self.regression.predict(X)
|
|||
|
|
|
|||
|
|
def evaluate(self, X: Sequence[Sequence[float]],
|
|||
|
|
y: Sequence[float]) -> Evaluation:
|
|||
|
|
pred = self.predict(X)
|
|||
|
|
return Evaluation(
|
|||
|
|
r2=r2_score(y, pred), mae=mae_score(y, pred),
|
|||
|
|
rmse=rmse_score(y, pred), n_samples=len(y))
|
|||
|
|
|
|||
|
|
def explain(self) -> List[Dict[str, Any]]:
|
|||
|
|
"""特征贡献度（尺度归一化：|权重| × 特征标准差），供可解释性。
|
|||
|
|
|
|||
|
|
归一化动机：原始权重受特征量纲影响（温度 850℃ vs 配比 30），
|
|||
|
|
``|权重|×std`` 才反映特征对预测的实际扰动幅度（与 sklearn
|
|||
|
|
permutation importance / 标准化系数同思路），可跨特征横向比较。
|
|||
|
|
"""
|
|||
|
|
if self.regression is None:
|
|||
|
|
raise ModelError("模型未训练")
|
|||
|
|
names = self.recipe.feature_names or [
|
|||
|
|
f"x{i}" for i in range(len(self.regression.coef_))]
|
|||
|
|
weights = list(self.regression.coef_)
|
|||
|
|
stds = getattr(self, "_feature_std", None) or [1.0] * len(weights)
|
|||
|
|
contribs = [abs(weights[i]) * stds[i] for i in range(len(weights))]
|
|||
|
|
total = sum(contribs) or 1.0
|
|||
|
|
return [{"feature": names[i], "weight": round(weights[i], 6),
|
|||
|
|
"importance": round(contribs[i] / total, 4)}
|
|||
|
|
for i in range(len(weights))]
|
|||
|
|
|
|||
|
|
def to_dict(self) -> Dict[str, Any]:
|
|||
|
|
return {
|
|||
|
|
"recipe": {"target": self.recipe.target, "alpha": self.recipe.alpha,
|
|||
|
|
"feature_names": list(self.recipe.feature_names),
|
|||
|
|
"unit": self.recipe.unit},
|
|||
|
|
"regression": (self.regression.to_dict()
|
|||
|
|
if self.regression else None),
|
|||
|
|
"feature_std": list(self._feature_std),
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_dict(cls, d: Dict[str, Any]) -> "QualityModel":
|
|||
|
|
model = cls(ModelRecipe.from_dict(d.get("recipe", {})))
|
|||
|
|
reg = d.get("regression")
|
|||
|
|
if reg:
|
|||
|
|
model.regression = RidgeRegression.from_dict(reg)
|
|||
|
|
model._feature_std = [float(x) for x in d.get("feature_std", [])]
|
|||
|
|
return model
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 训练集数据集（特征矩阵 + 标签）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass
|
|||
|
|
class TrainingSet:
|
|||
|
|
"""训练集：特征名 + X + y。"""
|
|||
|
|
|
|||
|
|
feature_names: List[str]
|
|||
|
|
X: List[List[float]]
|
|||
|
|
y: List[float]
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_records(cls, records: Sequence[Dict[str, float]],
|
|||
|
|
feature_names: Sequence[str],
|
|||
|
|
target: str) -> "TrainingSet":
|
|||
|
|
X, y = [], []
|
|||
|
|
for r in records:
|
|||
|
|
if target not in r or math.isnan(r[target]):
|
|||
|
|
continue
|
|||
|
|
row = [r.get(fn, NAN) for fn in feature_names]
|
|||
|
|
if any(math.isnan(v) for v in row):
|
|||
|
|
continue
|
|||
|
|
X.append(row)
|
|||
|
|
y.append(r[target])
|
|||
|
|
return cls(feature_names=list(feature_names), X=X, y=y)
|
|||
|
|
|
|||
|
|
def __len__(self) -> int:
|
|||
|
|
return len(self.X)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 超参包加载（零依赖 YAML 子集 / JSON）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_recipe(text: str) -> ModelRecipe:
|
|||
|
|
text = text.strip()
|
|||
|
|
if text.startswith("{"):
|
|||
|
|
data = json.loads(text)
|
|||
|
|
else:
|
|||
|
|
data = _parse_yaml_subset(text)
|
|||
|
|
if not isinstance(data, dict):
|
|||
|
|
raise ModelError("超参包顶层应为映射")
|
|||
|
|
return ModelRecipe.from_dict(data)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _parse_yaml_subset(text: str) -> Any:
|
|||
|
|
"""极简 YAML 子集解析（与 features.py 同款实现，避免跨模块依赖）。"""
|
|||
|
|
lines: List[str] = []
|
|||
|
|
for raw in text.splitlines():
|
|||
|
|
stripped = raw.rstrip()
|
|||
|
|
if not stripped.strip() or stripped.lstrip().startswith("#"):
|
|||
|
|
continue
|
|||
|
|
hi = _find_inline_comment(stripped)
|
|||
|
|
if hi is not None:
|
|||
|
|
stripped = stripped[:hi].rstrip()
|
|||
|
|
if stripped:
|
|||
|
|
lines.append(stripped)
|
|||
|
|
parser = _YamlParser(lines)
|
|||
|
|
return parser.parse_block(0)[0] if lines else {}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _find_inline_comment(line: str) -> Optional[int]:
|
|||
|
|
depth = 0
|
|||
|
|
in_str = False
|
|||
|
|
for i, ch in enumerate(line):
|
|||
|
|
if ch == '"':
|
|||
|
|
in_str = not in_str
|
|||
|
|
elif not in_str:
|
|||
|
|
if ch in "[{":
|
|||
|
|
depth += 1
|
|||
|
|
elif ch in "]}":
|
|||
|
|
depth = max(0, depth - 1)
|
|||
|
|
elif ch == "#" and depth == 0:
|
|||
|
|
if i == 0 or line[i - 1] in (" ", "\t"):
|
|||
|
|
return i
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _parse_scalar(raw: str) -> Any:
|
|||
|
|
raw = raw.strip()
|
|||
|
|
if not raw:
|
|||
|
|
return None
|
|||
|
|
if raw.startswith('"') and raw.endswith('"'):
|
|||
|
|
return raw[1:-1]
|
|||
|
|
if raw.startswith("[") or raw.startswith("{"):
|
|||
|
|
try:
|
|||
|
|
return json.loads(raw)
|
|||
|
|
except Exception:
|
|||
|
|
return raw
|
|||
|
|
low = raw.lower()
|
|||
|
|
if low == "true":
|
|||
|
|
return True
|
|||
|
|
if low == "false":
|
|||
|
|
return False
|
|||
|
|
if low in ("null", "~", "none"):
|
|||
|
|
return None
|
|||
|
|
try:
|
|||
|
|
return int(raw)
|
|||
|
|
except ValueError:
|
|||
|
|
pass
|
|||
|
|
try:
|
|||
|
|
return float(raw)
|
|||
|
|
except ValueError:
|
|||
|
|
pass
|
|||
|
|
return raw
|
|||
|
|
|
|||
|
|
|
|||
|
|
class _YamlParser:
|
|||
|
|
def __init__(self, lines: List[str]) -> None:
|
|||
|
|
self.lines = lines
|
|||
|
|
self.i = 0
|
|||
|
|
|
|||
|
|
def _indent(self, line: str) -> int:
|
|||
|
|
return len(line) - len(line.lstrip(" "))
|
|||
|
|
|
|||
|
|
def parse_block(self, indent: int) -> Tuple[Any, bool]:
|
|||
|
|
if self.i >= len(self.lines):
|
|||
|
|
return {}, False
|
|||
|
|
line = self.lines[self.i]
|
|||
|
|
cur = self._indent(line)
|
|||
|
|
if cur < indent:
|
|||
|
|
return {}, False
|
|||
|
|
stripped = line.strip()
|
|||
|
|
if stripped.startswith("- ") or stripped == "-":
|
|||
|
|
return self._parse_list(cur), True
|
|||
|
|
return self._parse_mapping(cur), False
|
|||
|
|
|
|||
|
|
def _parse_mapping(self, indent: int) -> Dict[str, Any]:
|
|||
|
|
result: Dict[str, Any] = {}
|
|||
|
|
effective = indent
|
|||
|
|
if self.i < len(self.lines):
|
|||
|
|
first = self._indent(self.lines[self.i])
|
|||
|
|
if first > indent:
|
|||
|
|
effective = first
|
|||
|
|
while self.i < len(self.lines):
|
|||
|
|
line = self.lines[self.i]
|
|||
|
|
cur = self._indent(line)
|
|||
|
|
if cur < effective:
|
|||
|
|
break
|
|||
|
|
if cur > effective:
|
|||
|
|
self.i += 1
|
|||
|
|
continue
|
|||
|
|
stripped = line.strip()
|
|||
|
|
if stripped.startswith("- "):
|
|||
|
|
break
|
|||
|
|
key, sep, rest = stripped.partition(":")
|
|||
|
|
if not sep:
|
|||
|
|
self.i += 1
|
|||
|
|
continue
|
|||
|
|
key = key.strip()
|
|||
|
|
rest = rest.strip()
|
|||
|
|
self.i += 1
|
|||
|
|
if rest:
|
|||
|
|
result[key] = _parse_scalar(rest)
|
|||
|
|
else:
|
|||
|
|
if self.i < len(self.lines):
|
|||
|
|
nxt = self._indent(self.lines[self.i])
|
|||
|
|
if nxt > effective:
|
|||
|
|
val, _ = self.parse_block(nxt)
|
|||
|
|
result[key] = val
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def _parse_list(self, indent: int) -> List[Any]:
|
|||
|
|
items: List[Any] = []
|
|||
|
|
while self.i < len(self.lines):
|
|||
|
|
line = self.lines[self.i]
|
|||
|
|
cur = self._indent(line)
|
|||
|
|
if cur < indent:
|
|||
|
|
break
|
|||
|
|
if cur > indent:
|
|||
|
|
self.i += 1
|
|||
|
|
continue
|
|||
|
|
stripped = line.strip()
|
|||
|
|
if not stripped.startswith("-"):
|
|||
|
|
break
|
|||
|
|
item_text = stripped[1:].strip()
|
|||
|
|
if not item_text:
|
|||
|
|
self.i += 1
|
|||
|
|
items.append(None)
|
|||
|
|
continue
|
|||
|
|
if ":" in item_text and not item_text.startswith('"'):
|
|||
|
|
k, sep, v = item_text.partition(":")
|
|||
|
|
if sep:
|
|||
|
|
item: Dict[str, Any] = {k.strip(): _parse_scalar(v.strip())}
|
|||
|
|
self.i += 1
|
|||
|
|
if self.i < len(self.lines):
|
|||
|
|
child_indent = self._indent(self.lines[self.i])
|
|||
|
|
if child_indent > cur:
|
|||
|
|
sub, _ = self.parse_block(child_indent)
|
|||
|
|
if isinstance(sub, dict):
|
|||
|
|
item.update(sub)
|
|||
|
|
items.append(item)
|
|||
|
|
continue
|
|||
|
|
items.append(_parse_scalar(item_text))
|
|||
|
|
self.i += 1
|
|||
|
|
return items
|