our-claude-skills/custom-skills/32-notion-writer/code/scripts/ntn_files.py

#!/usr/bin/env python3
"""Local file uploads to Notion via the `ntn` CLI.

Owns all subprocess I/O for the skill. The CLI handles the full File Upload
lifecycle (create -> send bytes -> complete) and prints the upload ID.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional, Callable, Any

_WORKSPACE: Optional[Dict[str, str]] = None


def _ntn_env():
    """Env for ntn subprocesses: force ntn to authenticate as the SAME
    integration the script uses (NOTION_API_KEY/NOTION_TOKEN), so an uploaded
    file and the page that references it share one identity. Notion scopes
    file uploads to the creating integration, so a mismatch makes the upload
    un-attachable."""
    env = dict(os.environ)
    token = os.environ.get("NOTION_API_KEY") or os.environ.get("NOTION_TOKEN")
    if token:
        env["NOTION_API_TOKEN"] = token
    return env


class NtnUploadError(Exception):
    """Raised when `ntn` is unavailable or a file upload fails."""

    def __init__(self, message: str, path: str = "", stderr: str = ""):
        super().__init__(message)
        self.path = path
        self.stderr = stderr


def preflight() -> Dict[str, str]:
    """Verify `ntn` is installed and logged in; return its target workspace.

    Cached for the process. Raises NtnUploadError with an actionable hint
    on failure. Prints one informational line naming the workspace `ntn`
    targets (file uploads are workspace-scoped).
    """
    global _WORKSPACE
    if _WORKSPACE is not None:
        return _WORKSPACE

    if shutil.which("ntn") is None:
        raise NtnUploadError(
            "ntn CLI not found. Install it with: "
            "curl -fsSL https://ntn.dev | bash"
        )

    result = subprocess.run(
        ["ntn", "api", "v1/users/me"],
        capture_output=True, text=True, env=_ntn_env(),
    )
    if result.returncode != 0:
        raise NtnUploadError(
            "ntn is not logged in. Run: ntn login\n" + result.stderr.strip()
        )

    try:
        me = json.loads(result.stdout)
        bot = me.get("bot", {})
        info = {
            "workspace_name": bot.get("workspace_name", "unknown"),
            "workspace_id": bot.get("workspace_id", "unknown"),
        }
    except (ValueError, AttributeError):
        info = {"workspace_name": "unknown", "workspace_id": "unknown"}

    print(f'ntn -> workspace "{info["workspace_name"]}"', file=sys.stderr)
    _WORKSPACE = info
    return info


def upload(path: Path) -> str:
    """Upload a local file via `ntn files create --plain`; return its id."""
    path = Path(path)
    with open(path, "rb") as fh:
        result = subprocess.run(
            ["ntn", "files", "create", "--plain"],
            stdin=fh, capture_output=True, text=True, env=_ntn_env(),
        )
    if result.returncode != 0:
        raise NtnUploadError(
            f"ntn upload failed for {path}", path=str(path),
            stderr=result.stderr.strip(),
        )
    if not result.stdout.strip():
        raise NtnUploadError(
            "ntn returned no output", path=str(path), stderr=result.stderr.strip()
        )
    first_line = result.stdout.strip().splitlines()[0]
    upload_id = first_line.split("\t")[0].strip()
    return upload_id


_REMOTE_PREFIXES = ("http://", "https://")


def _resolve_local(url: str, base_dir: Path) -> Optional[Path]:
    """Return the resolved path for a local media url, or None if remote."""
    if url.startswith(_REMOTE_PREFIXES):
        return None
    p = Path(url)
    return p if p.is_absolute() else Path(base_dir) / p


def materialize_local_media(
    blocks: List[Dict[str, Any]],
    base_dir: Path,
    upload_fn: Callable[[Path], str] = upload,
) -> List[Dict[str, Any]]:
    """Walk blocks (and nested children); upload local image files and
    rewrite them to file_upload shape. Remote images are left as-is."""
    for block in blocks:
        if block.get("type") == "image":
            img = block["image"]
            if img.get("type") == "external":
                url = img.get("external", {}).get("url", "")
                resolved = _resolve_local(url, base_dir)
                if resolved is not None:
                    if not resolved.is_file():
                        raise NtnUploadError(
                            f"image references missing file: {url}",
                            path=str(resolved),
                        )
                    upload_id = upload_fn(resolved)
                    caption = img.get("caption")
                    new_img: Dict[str, Any] = {
                        "type": "file_upload",
                        "file_upload": {"id": upload_id},
                    }
                    if caption:
                        new_img["caption"] = caption
                    block["image"] = new_img
        # Recurse into any nested children list.
        btype = block.get("type")
        nested = block.get(btype, {}) if isinstance(block.get(btype), dict) else {}
        if isinstance(nested.get("children"), list):
            materialize_local_media(nested["children"], base_dir, upload_fn)
        if btype == "column_list":
            for col in nested.get("children", []):
                col_children = col.get("column", {}).get("children", [])
                materialize_local_media(col_children, base_dir, upload_fn)
    return blocks