InvenTree：使用 Python API 脚本批量创建库位

#!/usr/bin/env python3
"""在 InvenTree 中的 Apothekerschrank 库位下创建 Schublade 子库位。

本脚本通过 REST API 连接到 InvenTree 实例，在一个名为
"Apothekerschrank" 的父库位下创建一组抽屉（"Schublade"）子库位。
对于 A、B、C 三行中的每一行，创建编号为 1 到 19 的抽屉，
共计 57 个子库位。每个子库位都会被分配 "Schublade" 库位类型。

本脚本是幂等的：如果父库位下已存在同名子库位，则会跳过
而不是重复创建。

配置从同目录下的 ``config.yaml`` 读取：

.. code-block:: yaml

    inventree:
      server: https://inventree.example.com
      token: your-api-token-here

依赖：
    - Python 3.8+
    - requests
    - PyYAML

用法::

    python3 create_locations.py
"""

import sys
import requests
import yaml
from pathlib import Path

# YAML 配置文件路径，应与本脚本位于同一目录。
CONFIG_PATH = Path(__file__).parent / "config.yaml"


def load_config():
    """从 ``config.yaml`` 加载 InvenTree 连接设置。

    返回包含 ``server`` 和 ``token`` 键的字典。
    """
    with open(CONFIG_PATH, "r") as f:
        return yaml.safe_load(f)["inventree"]


class InvenTreeAPI:
    """InvenTree 的最小化 REST API 客户端。

    封装 :mod:`requests`，提供基于 token 的认证以及
    针对 InvenTree API 的 GET 和 POST 便捷方法。
    """

    def __init__(self, server, token):
        """初始化 API 客户端。

        :param server: InvenTree 实例的基础 URL，例如
            ``https://inventree.example.com``。
        :param token: API 认证 token（可在 InvenTree Web 界面的
            *Settings > API Keys* 中生成）。
        """
        self.server = server.rstrip("/")
        self.token = token
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Token {token}",
            "Content-Type": "application/json",
        })

    def get(self, path, params=None):
        """执行 GET 请求并返回解析后的 JSON 响应。"""
        r = self.session.get(f"{self.server}{path}", params=params)
        r.raise_for_status()
        return r.json()

    def post(self, path, data):
        """执行带 JSON body 的 POST 请求并返回解析后的响应。"""
        r = self.session.post(f"{self.server}{path}", json=data)
        if r.status_code == 400:
            print(f"  错误 400：{r.json()}")
        r.raise_for_status()
        return r.json()


def find_location_type(api, name):
    """按名称查找库位类型。

    通过分页 API 调用遍历所有库位类型。
    找到则返回主键（``pk``），否则返回 ``None``。
    """
    offset = 0
    while True:
        data = api.get("/api/stock/location-type/", params={"limit": 100, "offset": offset})
        results = data if isinstance(data, list) else data.get("results", [])
        for t in results:
            if t["name"] == name:
                return t["pk"]
        if isinstance(data, list) or not data.get("next"):
            break
        offset += 100
    return None


def find_location(api, name):
    """在所有库位中按名称查找库位。

    通过分页 API 调用遍历所有库位。
    找到则返回主键（``pk``），否则返回 ``None``。
    """
    offset = 0
    while True:
        data = api.get("/api/stock/location/", params={"limit": 100, "offset": offset})
        for loc in data["results"]:
            if loc["name"] == name:
                return loc["pk"]
        if not data["next"]:
            break
        offset += 100
    return None


def find_child_location(api, parent_pk, name):
    """在指定父库位下按名称查找子库位。

    查询 API 中所有 ``parent`` 匹配 ``parent_pk`` 的库位，
    返回第一个匹配项的 ``pk``，否则返回 ``None``。
    """
    results = api.get("/api/stock/location/", params={
        "parent": parent_pk, "limit": 1000
    })
    loc_list = results if isinstance(results, list) else results.get("results", [])
    for loc in loc_list:
        if loc["name"] == name:
            return loc["pk"]
    return None


def main():
    """主入口：查找父库位、确保库位类型存在、创建子库位。"""
    config = load_config()
    api = InvenTreeAPI(config["server"], config["token"])

    # --- 查找或创建 "Schublade" 库位类型 ---
    # InvenTree 支持用户自定义库位类型，可分配给各个库位。
    # 如果该类型已存在则直接复用。
    schublade_type_pk = find_location_type(api, "Schublade")
    if schublade_type_pk:
        print(f"已找到库位类型 'Schublade'（pk={schublade_type_pk}）")
    else:
        result = api.post("/api/stock/location-type/", {
            "name": "Schublade",
            "description": "Schublade im Apothekerschrank",
            "icon": "ti:box:outline",
        })
        schublade_type_pk = result["pk"]
        print(f"已创建库位类型 'Schublade'（pk={schublade_type_pk}）")

    # --- 查找 "Apothekerschrank" 父库位 ---
    # 该库位必须已在 InvenTree 中存在。如果找不到，
    # 脚本将以错误退出。
    apotheker_pk = find_location(api, "Apothekerschrank")
    if not apotheker_pk:
        print("错误：未找到库位 'Apothekerschrank'")
        sys.exit(1)
    print(f"已找到库位 'Apothekerschrank'（pk={apotheker_pk}）")

    # --- 创建子库位：Schublade A1..A19、B1..B19、C1..C19 ---
    # 每个子库位都作为 Apothekerschrank 库位的子项创建，
    # 并分配 "Schublade" 库位类型。已存在的库位会被跳过，
    # 以保持脚本的幂等性。
    created = 0
    skipped = 0
    for letter in ["A", "B", "C"]:
        for number in range(1, 20):
            name = f"Schublade {letter}{number}"
            existing = find_child_location(api, apotheker_pk, name)
            if existing:
                print(f"  跳过（已存在）：{name}（pk={existing}）")
                skipped += 1
                continue
            result = api.post("/api/stock/location/", {
                "name": name,
                "parent": apotheker_pk,
                "location_type": schublade_type_pk,
            })
            print(f"  已创建：{name}（pk={result['pk']}）")
            created += 1

    print(f"\n完成：已创建 {created} 个，跳过 {skipped} 个")


if __name__ == "__main__":
    main()