InvenTree: масовий експорт PDF-міток місць зберігання через Python-скрипт

#!/usr/bin/env python3
"""Завантаження та об'єднання PDF-міток місць зберігання з InvenTree.

Цей скрипт підключається до інстансу InvenTree через REST API, отримує
місця зберігання (за потреби відфільтровані за шаблоном назви та/або
батьківським місцем), генерує PDF-мітки паралельно за допомогою шаблона
міток InvenTree та об'єднує результати в один PDF-файл.

Конфігурація зчитується з ``config.yaml`` у тому ж каталозі::

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

Вимоги:
    - Python 3.8+
    - requests
    - PyYAML
    - pypdf

Приклади використання::

    # Згенерувати мітки для УСІХ місць зберігання, об'єднані в один PDF
    python3 export_location_labels.py

    # Фільтр за назвою: glob-шаблон (з *) або підрядок (без *)
    python3 export_location_labels.py -q "Schublade A*"
    python3 export_location_labels.py -q "Schublade"

    # Фільтр за батьківським місцем (назва або числовий pk)
    python3 export_location_labels.py -p Apothekerschrank
    python3 export_location_labels.py -p 9

    # Комбінування обох фільтрів
    python3 export_location_labels.py -q "Schublade A*" -p Apothekerschrank

    # Виключити конкретні місця за назвою (glob або підрядок, ті ж правила, що й -q)
    python3 export_location_labels.py -e "Test*" -e "Illerbeuren"

    # Включити лише конкретні місця (переважає над виключеннями та -q)
    python3 export_location_labels.py -i "Schublade A1" -i "Schublade B2"

    # Включення + виключення разом: включення має пріоритет
    python3 export_location_labels.py -i "Schublade*" -e "Schublade C*"

    # Вибрати конкретний шаблон мітки (-t скорочення)
    python3 export_location_labels.py -t "Lagerort Groß 62mm"

    # Власний шлях виводу для об'єднаного PDF
    python3 export_location_labels.py -o ./my_labels.pdf

    # Запис окремих PDF у каталог замість об'єднання
    python3 export_location_labels.py --individual -o ./my_labels_dir/

    # Автоматичне ім'я виводу: ім'я файлу походить від пошукового терміна
    # "Schublade A*" -> Schublade_A.pdf (об'єднаний) або Schublade_A/ (окремі)
    python3 export_location_labels.py -q "Schublade A*"

    # Налаштування паралелізму та таймауту
    python3 export_location_labels.py --workers 16 --timeout 120

Фільтрування:
    Усі зіставлення назв нормалізують пробіли: будь-яка послідовність
    пробільних символів (пробіли, табуляції, символи нового рядка тощо)
    у назві місця та в шаблоні query/exclude/include згортається до
    одного пробілу перед порівнянням.

    - **Фільтр назви** (``-q``): Якщо рядок запиту містить ``*``, він
      трактується як glob-шаблон (наприклад ``"Schublade A*"`` відповідає
      ``Schublade A1``, ``Schublade A10`` тощо). Якщо ``*`` відсутній,
      використовується регістронезалежне зіставлення підрядка.

    - **Фільтр батьківського місця** (``-p``): Фільтрує місця, чий
      батьківський елемент відповідає заданому значенню. Значення може
      бути або назвою батьківського місця (наприклад ``Apothekerschrank``),
      або його числовим первинним ключем (наприклад ``9``).

    - **Виключення** (``-e``): Виключити місця, що відповідають заданому
      шаблону. Можна вказувати кілька разів. Ті ж правила glob/підрядок,
      що й для ``-q``. Місце виключається, якщо воно відповідає *будь-якому*
      шаблону виключення.

    - **Включення** (``-i``): Включити лише місця, що відповідають
      заданому шаблону. Можна вказувати кілька разів. Ті ж правила
      glob/підрядок, що й для ``-q``. Місце включається, якщо воно
      відповідає *будь-якому* шаблону включення. **Включення переважає
      над як ``-q``, так і ``-e``**: коли задано ``-i``, фільтр запиту
      ігнорується, а виключені місця, що відповідають шаблону включення,
      все одно включаються.

Вивід:
    За замовчуванням усі згенеровані PDF-мітки об'єднуються в один
    PDF-файл за допомогою ``pypdf``. Коли задано ``--individual``, кожна
    мітка записується як окремий PDF-файл у каталозі.

    Шлях виводу визначається так:

    1. Якщо задано ``-o``, він використовується безпосередньо (шлях до
       файлу для режиму об'єднання, шлях до каталогу для режиму окремих
       файлів).
    2. Якщо ``-o`` не задано, ім'я виводиться автоматично:
       - З пошукового терміна ``-q`` із видаленням glob-символів
         (``*``, ``?``) і заміною пробілів на підкреслення, наприклад
         ``"Schublade A*"`` -> ``Schublade_A.pdf``.
       - Якщо немає ``-q``, але є ``-i`` — з першого шаблону включення
         (те ж видалення).
       - Якщо ні того, ні іншого — ``StockLocationLabels.pdf``
         (об'єднаний) або ``StockLocationLabels/`` (окремі).

Як це працює:
    1. Отримує всі місця зберігання з API InvenTree (пагіновано).
    2. Застосовує опціональні фільтри назви, батьківського місця,
       включення та виключення.
    3. Отримує доступні шаблони міток ``stocklocation``.
    4. Паралельно подає завдання друку мітки для кожного місця, що
       відповідає, через ``POST /api/label/print/``.
    5. Опитує ``GET /api/data-output/<pk>/``, поки кожне завдання не
       завершиться.
    6. Завантажує згенерований PDF із шляху виводу.
    7. Об'єднує всі окремі PDF в один за допомогою ``pypdf`` (або
       записує їх окремо, якщо задано ``--individual``).
"""

import argparse
import fnmatch
import io
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import requests
import yaml
from pypdf import PdfWriter, PdfReader

CONFIG_PATH = Path(__file__).parent / "config.yaml"


def load_config():
    with open(CONFIG_PATH, "r") as f:
        return yaml.safe_load(f)["inventree"]


class InvenTreeAPI:
    def __init__(self, server, token):
        self.server = server.rstrip("/")
        self.token = token
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Token {token}",
        })

    def get(self, path, params=None):
        r = self.session.get(f"{self.server}{path}", params=params)
        r.raise_for_status()
        return r

    def post(self, path, data=None, json=None):
        r = self.session.post(f"{self.server}{path}", data=data, json=json)
        if r.status_code == 400:
            print(f"  ERROR 400: {r.text}")
        r.raise_for_status()
        return r


def get_all_locations(api):
    """Отримати всі місця зберігання через пагіновані виклики API."""
    locations = []
    offset = 0
    while True:
        r = api.get("/api/stock/location/", params={
            "limit": 100, "offset": offset,
        })
        data = r.json()
        locations.extend(data["results"])
        if not data["next"]:
            break
        offset += 100
    return locations


def get_location_templates(api):
    """Отримати всі увімкнені шаблони міток для місць зберігання."""
    templates = []
    offset = 0
    while True:
        r = api.get("/api/label/template/", params={
            "limit": 100, "offset": offset,
            "model_type": "stocklocation", "enabled": True,
        })
        data = r.json()
        templates.extend(data["results"])
        if not data["next"]:
            break
        offset += 100
    return templates


def print_and_download_label(api, template_pk, item_pks, timeout=60):
    """Подати завдання друку мітки, опитувати до завершення, завантажити PDF.

    Повертає вміст PDF у вигляді байтів.
    """
    r = api.post("/api/label/print/", json={
        "template": template_pk,
        "items": item_pks,
    })
    result = r.json()
    output_pk = result["pk"]

    deadline = time.time() + timeout
    while time.time() < deadline:
        r = api.get(f"/api/data-output/{output_pk}/")
        data = r.json()
        if data.get("complete"):
            output_path = data.get("output")
            if not output_path:
                raise RuntimeError(
                    f"Вивід мітки {output_pk} завершено, але немає шляху виводу"
                )
            pdf_url = f"{api.server}{output_path}"
            pr = api.session.get(pdf_url)
            pr.raise_for_status()
            return pr.content
        time.sleep(0.5)

    raise TimeoutError(
        f"Вивід мітки {output_pk} не завершено протягом {timeout}с"
    )


def sanitize_filename(name):
    """Зробити рядок безпечним для використання як ім'я файлу."""
    for ch in r'<>:"/\\|?*':
        name = name.replace(ch, "_")
    return name.strip()


def normalize_ws(s):
    """Згорнути всі послідовності пробільних символів у рядку до одиночних пробілів."""
    return re.sub(r"\s+", " ", s).strip()


def name_matches(name, pattern):
    """Перевірити, чи відповідає назва місця шаблону.

    Glob, якщо шаблон містить *, інакше регістронезалежний підрядок.
    Пробіли нормалізуються в обох перед порівнянням.
    """
    name = normalize_ws(name)
    pattern = normalize_ws(pattern)
    if "*" in pattern:
        return fnmatch.fnmatch(name, pattern)
    return pattern.lower() in name.lower()


def filter_locations(locations, query, parent, includes, excludes):
    """Відфільтрувати місця за шаблоном назви, батьківським місцем, списками включення/виключення.

    - query: glob-шаблон (якщо містить *) або регістронезалежний підрядок
    - parent: назва батьківського місця або pk (рядок int)
    - includes: список шаблонів; якщо непорожній, залишаються лише
      місця, що відповідають (переважає над query та excludes)
    - excludes: список шаблонів; місця, що відповідають, вилучаються
    """
    if includes:
        filtered = [l for l in locations if any(name_matches(l.get("name", ""), p) for p in includes)]
    else:
        filtered = locations

        if query:
            filtered = [l for l in filtered if name_matches(l.get("name", ""), query)]

    if parent:
        parent_pk = None
        if parent.isdigit():
            parent_pk = int(parent)
        else:
            for l in locations:
                if normalize_ws(l.get("name", "")) == normalize_ws(parent):
                    parent_pk = l["pk"]
                    break
            if parent_pk is None:
                print(f"ПОМИЛКА: Батьківське місце '{parent}' не знайдено")
                sys.exit(1)
        filtered = [l for l in filtered if l.get("parent") == parent_pk]

    if excludes and not includes:
        filtered = [l for l in filtered if not any(name_matches(l.get("name", ""), p) for p in excludes)]
    elif excludes and includes:
        filtered = [l for l in filtered if not any(name_matches(l.get("name", ""), p) for p in excludes) or any(name_matches(l.get("name", ""), p) for p in includes)]

    return filtered


def main():
    parser = argparse.ArgumentParser(
        description="Завантаження та об'єднання PDF-міток місць зберігання з InvenTree"
    )
    parser.add_argument(
        "-q", "--query", default=None,
        help="Фільтрувати місця за назвою (glob, якщо містить *, інакше підрядок)",
    )
    parser.add_argument(
        "-p", "--parent", default=None,
        help="Фільтр за батьківським місцем (назва або числовий pk)",
    )
    parser.add_argument(
        "-e", "--exclude", action="append", default=[],
        help="Виключити місця, що відповідають цьому шаблону (glob або підрядок). Можна вказувати кілька разів.",
    )
    parser.add_argument(
        "-i", "--include", action="append", default=[],
        help="Включити лише місця, що відповідають цьому шаблону (glob або підрядок). Переважає над -q та -e. Можна вказувати кілька разів.",
    )
    parser.add_argument(
        "-t", "--template", default=None,
        help="Назва шаблону мітки для використання (за замовчуванням: перший доступний)",
    )
    parser.add_argument(
        "-o", "--output", default=None,
        help="Шлях виводу: PDF-файл (режим об'єднання) або каталог (режим окремих файлів). "
             "Якщо не задано, виводиться автоматично з пошукового терміна або шаблону включення.",
    )
    parser.add_argument(
        "--individual", action="store_true",
        help="Записати окремі PDF у каталог замість об'єднання в один PDF",
    )
    parser.add_argument(
        "--workers", type=int, default=8,
        help="Кількість паралельних завдань друку (за замовчуванням: 8)",
    )
    parser.add_argument(
        "--timeout", type=int, default=60,
        help="Таймаут у секундах на завдання мітки (за замовчуванням: 60)",
    )
    args = parser.parse_args()

    config = load_config()
    api = InvenTreeAPI(config["server"], config["token"])

    # --- Отримати шаблони міток ---
    templates = get_location_templates(api)
    if not templates:
        print("ПОМИЛКА: Не знайдено увімкнених шаблонів міток для model_type 'stocklocation'")
        sys.exit(1)

    print(f"Знайдено {len(templates)} шаблон(ів) міток місць зберігання:")
    for t in templates:
        print(f"  - {t['name']} (pk={t['pk']}, {t['width']}x{t['height']}мм)")

    selected = None
    if args.template:
        for t in templates:
            if t["name"] == args.template:
                selected = t
                break
        if not selected:
            print(f"ПОМИЛКА: Шаблон '{args.template}' не знайдено")
            sys.exit(1)
    else:
        selected = templates[0]
    print(f"\nВикористовується шаблон: {selected['name']} (pk={selected['pk']})")

    # --- Отримати та відфільтрувати місця ---
    locations = get_all_locations(api)
    print(f"Знайдено {len(locations)} місць зберігання загалом")

    locations = filter_locations(locations, args.query, args.parent, args.include, args.exclude)
    print(f"Після фільтрування: {len(locations)} місць")

    if not locations:
        print("Жодне місце не відповідає критеріям фільтру.")
        return

    for loc in locations:
        print(f"  {loc['name']} (pk={loc['pk']})")

    # --- Друк міток паралельно ---
    print(f"\nГенерування {len(locations)} міток паралельно "
          f"({args.workers} воркерів)...")

    results = {}  # pk -> (назва, pdf_bytes або None)
    errors = {}

    def _print_one(loc):
        name = loc.get("name", f"location_{loc['pk']}")
        pk = loc["pk"]
        try:
            pdf = print_and_download_label(
                api, selected["pk"], [pk], timeout=args.timeout
            )
            return pk, name, pdf, None
        except Exception as e:
            return pk, name, None, str(e)

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(_print_one, loc): loc for loc in locations}
        for fut in as_completed(futures):
            pk, name, pdf, err = fut.result()
            if err:
                print(f"  НЕВДАЧА: {name} (pk={pk}): {err}")
                errors[pk] = err
            else:
                print(f"  OK: {name} (pk={pk}, {len(pdf)} байт)")
                results[pk] = (name, pdf)

    if not results:
        print("\nПОМИЛКА: Жодну мітку не вдалося згенерувати.")
        sys.exit(1)

    # --- Визначити шлях виводу ---
    def derive_name():
        """Автоматично вивести ім'я виводу з query або першого шаблону включення."""
        source = None
        if args.query:
            source = args.query
        elif args.include:
            source = args.include[0]
        if source:
            # Видалити glob-символи, нормалізувати пробіли, замінити пробіли на _
            cleaned = re.sub(r"[*?]", "", source)
            cleaned = normalize_ws(cleaned).replace(" ", "_")
            cleaned = sanitize_filename(cleaned)
            return cleaned if cleaned else "StockLocationLabels"
        return "StockLocationLabels"

    if args.output:
        output_path = Path(args.output)
    else:
        base_name = derive_name()
        if args.individual:
            output_path = Path(base_name)
        else:
            output_path = Path(f"{base_name}.pdf")

    # --- Запис виводу ---
    if args.individual:
        output_path.mkdir(parents=True, exist_ok=True)
        print(f"\nЗапис {len(results)} окремих PDF у {output_path}/...")
        for pk in sorted(results.keys()):
            name, pdf_bytes = results[pk]
            safe_name = sanitize_filename(normalize_ws(name).replace(" ", "_"))
            pdf_path = output_path / f"{safe_name}.pdf"
            pdf_path.write_bytes(pdf_bytes)
            print(f"  {pdf_path.name} ({len(pdf_bytes)} байт)")
        print(f"\nГотово: {len(results)} міток записано, {len(errors)} невдач")
        print(f"Каталог виводу: {output_path.resolve()}")
    else:
        print(f"\nОб'єднання {len(results)} PDF у {output_path}...")
        writer = PdfWriter()
        for pk in sorted(results.keys()):
            name, pdf_bytes = results[pk]
            reader = PdfReader(io.BytesIO(pdf_bytes))
            for page in reader.pages:
                writer.add_page(page)

        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as f:
            writer.write(f)

        print(f"\nГотово: {len(results)} міток об'єднано, {len(errors)} невдач")
        print(f"Вивід: {output_path.resolve()}")


if __name__ == "__main__":
    main()