InvenTree bulk stock location label PDF export via Python scripting

#!/usr/bin/env python3
"""Download and merge PDF storage location labels from InvenTree.

This script connects to an InvenTree instance via the REST API, retrieves
stock locations (optionally filtered by name pattern and/or parent
location), generates PDF labels in parallel using an InvenTree label
template, and merges the results into a single PDF file.

Configuration is read from ``config.yaml`` in the same directory::

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

Requirements:
    - Python 3.8+
    - requests
    - PyYAML
    - pypdf

Usage examples::

    # Generate labels for ALL storage locations, merged into one PDF
    python3 export_location_labels.py

    # Filter by name: glob pattern (with *) or substring (without *)
    python3 export_location_labels.py -q "Schublade A*"
    python3 export_location_labels.py -q "Schublade"

    # Filter by parent location (name or numeric pk)
    python3 export_location_labels.py -p Apothekerschrank
    python3 export_location_labels.py -p 9

    # Combine both filters
    python3 export_location_labels.py -q "Schublade A*" -p Apothekerschrank

    # Exclude specific locations by name (glob or substring, same rules as -q)
    python3 export_location_labels.py -e "Test*" -e "Illerbeuren"

    # Include only specific locations (overrides excludes and -q)
    python3 export_location_labels.py -i "Schublade A1" -i "Schublade B2"

    # Include + exclude combined: include takes precedence
    python3 export_location_labels.py -i "Schublade*" -e "Schublade C*"

    # Select a specific label template (-t shorthand)
    python3 export_location_labels.py -t "Lagerort Groß 62mm"

    # Custom output path for merged PDF
    python3 export_location_labels.py -o ./my_labels.pdf

    # Write individual PDFs to a directory instead of merging
    python3 export_location_labels.py --individual -o ./my_labels_dir/

    # Auto-named output: derives filename from the search term
    # "Schublade A*" -> Schublade_A.pdf (merged) or Schublade_A/ (individual)
    python3 export_location_labels.py -q "Schublade A*"

    # Adjust parallelism and timeout
    python3 export_location_labels.py --workers 16 --timeout 120

Filtering:
    All name matching normalizes whitespace: any sequence of whitespace
    characters (spaces, tabs, newlines, etc.) in both the location name
    and the query/exclude/include pattern is collapsed to a single space
    before comparison.

    - **Name filter** (``-q``): If the query string contains ``*``, it is
      treated as a glob pattern (e.g. ``"Schublade A*"`` matches
      ``Schublade A1``, ``Schublade A10``, etc.).  If no ``*`` is present,
      a case-insensitive substring match is used.

    - **Parent filter** (``-p``): Filters locations whose parent matches
      the given value.  The value can be either the parent location name
      (e.g. ``Apothekerschrank``) or its numeric primary key (e.g. ``9``).

    - **Exclude** (``-e``): Exclude locations matching the given pattern.
      Can be specified multiple times.  Same glob/substring rules as ``-q``.
      A location is excluded if it matches *any* exclude pattern.

    - **Include** (``-i``): Include only locations matching the given
      pattern.  Can be specified multiple times.  Same glob/substring rules
      as ``-q``.  A location is included if it matches *any* include
      pattern.  **Include overrides both ``-q`` and ``-e``**: when
      ``-i`` is provided, the query filter is ignored and excluded
      locations that match an include pattern are still included.

Output:
    By default, all generated label PDFs are merged into a single PDF
    file using ``pypdf``.  When ``--individual`` is given, each label is
    written as a separate PDF file in a directory instead.

    The output path is determined as follows:

    1. If ``-o`` is given, it is used directly (file path for merged
       mode, directory path for individual mode).
    2. If ``-o`` is not given, the name is auto-derived:
       - From the ``-q`` search term with glob characters (``*``, ``?``)
         stripped and spaces replaced by underscores, e.g.
         ``"Schublade A*"`` -> ``Schublade_A.pdf``.
       - If no ``-q`` but ``-i`` is given, from the first include pattern
         (same stripping).
       - If neither, ``StockLocationLabels.pdf`` (merged) or
         ``StockLocationLabels/`` (individual).

How it works:
    1. Fetches all stock locations from the InvenTree API (paginated).
    2. Applies the optional name, parent, include and exclude filters.
    3. Fetches available ``stocklocation`` label templates.
    4. Submits a label print job for each matching location in parallel
       via ``POST /api/label/print/``.
    5. Polls ``GET /api/data-output/<pk>/`` until each job is complete.
    6. Downloads the generated PDF from the output path.
    7. Merges all individual PDFs into one using ``pypdf`` (or writes
       them individually if ``--individual`` is given).
"""

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):
    """Fetch all stock locations via paginated API calls."""
    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):
    """Fetch all enabled label templates for stock locations."""
    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):
    """Submit a label print job, poll until complete, download the PDF.

    Returns the PDF content as bytes.
    """
    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"Label output {output_pk} complete but no output path"
                )
            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"Label output {output_pk} did not complete within {timeout}s"
    )


def sanitize_filename(name):
    """Make a string safe for use as a filename."""
    for ch in r'<>:"/\\|?*':
        name = name.replace(ch, "_")
    return name.strip()


def normalize_ws(s):
    """Collapse all whitespace sequences in a string to single spaces."""
    return re.sub(r"\s+", " ", s).strip()


def name_matches(name, pattern):
    """Check if a location name matches a pattern.

    Glob if pattern contains *, otherwise case-insensitive substring.
    Whitespace is normalized in both before comparison.
    """
    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):
    """Filter locations by name pattern, parent, include/exclude lists.

    - query: glob pattern (if contains *) or case-insensitive substring
    - parent: parent location name or pk (int string)
    - includes: list of patterns; if non-empty, only matching locations
      are kept (overrides query and excludes)
    - excludes: list of patterns; matching locations are removed
    """
    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"ERROR: Parent location '{parent}' not found")
                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="Download and merge PDF storage location labels from InvenTree"
    )
    parser.add_argument(
        "-q", "--query", default=None,
        help="Filter locations by name (glob if contains *, else substring)",
    )
    parser.add_argument(
        "-p", "--parent", default=None,
        help="Filter by parent location (name or numeric pk)",
    )
    parser.add_argument(
        "-e", "--exclude", action="append", default=[],
        help="Exclude locations matching this pattern (glob or substring). Can be given multiple times.",
    )
    parser.add_argument(
        "-i", "--include", action="append", default=[],
        help="Include only locations matching this pattern (glob or substring). Overrides -q and -e. Can be given multiple times.",
    )
    parser.add_argument(
        "-t", "--template", default=None,
        help="Name of the label template to use (default: first available)",
    )
    parser.add_argument(
        "-o", "--output", default=None,
        help="Output path: PDF file (merged mode) or directory (individual mode). "
             "If not given, auto-derived from the search term or include pattern.",
    )
    parser.add_argument(
        "--individual", action="store_true",
        help="Write individual PDFs to a directory instead of merging into one PDF",
    )
    parser.add_argument(
        "--workers", type=int, default=8,
        help="Number of parallel print jobs (default: 8)",
    )
    parser.add_argument(
        "--timeout", type=int, default=60,
        help="Timeout in seconds per label job (default: 60)",
    )
    args = parser.parse_args()

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

    # --- Fetch label templates ---
    templates = get_location_templates(api)
    if not templates:
        print("ERROR: No enabled label templates found for model_type 'stocklocation'")
        sys.exit(1)

    print(f"Found {len(templates)} stock location label template(s):")
    for t in templates:
        print(f"  - {t['name']} (pk={t['pk']}, {t['width']}x{t['height']}mm)")

    selected = None
    if args.template:
        for t in templates:
            if t["name"] == args.template:
                selected = t
                break
        if not selected:
            print(f"ERROR: Template '{args.template}' not found")
            sys.exit(1)
    else:
        selected = templates[0]
    print(f"\nUsing template: {selected['name']} (pk={selected['pk']})")

    # --- Fetch and filter locations ---
    locations = get_all_locations(api)
    print(f"Found {len(locations)} total stock locations")

    locations = filter_locations(locations, args.query, args.parent, args.include, args.exclude)
    print(f"After filtering: {len(locations)} locations")

    if not locations:
        print("No locations match the filter criteria.")
        return

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

    # --- Print labels in parallel ---
    print(f"\nGenerating {len(locations)} labels in parallel "
          f"({args.workers} workers)...")

    results = {}  # pk -> (name, pdf_bytes or 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"  FAILED: {name} (pk={pk}): {err}")
                errors[pk] = err
            else:
                print(f"  OK: {name} (pk={pk}, {len(pdf)} bytes)")
                results[pk] = (name, pdf)

    if not results:
        print("\nERROR: No labels were generated successfully.")
        sys.exit(1)

    # --- Determine output path ---
    def derive_name():
        """Auto-derive output name from query or first include pattern."""
        source = None
        if args.query:
            source = args.query
        elif args.include:
            source = args.include[0]
        if source:
            # Strip glob characters, normalize whitespace, replace spaces with _
            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")

    # --- Write output ---
    if args.individual:
        output_path.mkdir(parents=True, exist_ok=True)
        print(f"\nWriting {len(results)} individual PDFs to {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)} bytes)")
        print(f"\nDone: {len(results)} labels written, {len(errors)} failed")
        print(f"Output directory: {output_path.resolve()}")
    else:
        print(f"\nMerging {len(results)} PDFs into {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"\nDone: {len(results)} labels merged, {len(errors)} failed")
        print(f"Output: {output_path.resolve()}")


if __name__ == "__main__":
    main()