folderlake 0.2.0__py3-none-any.whl

folderlake/extractors.py

@@ -0,0 +1,360 @@
+"""One small extractor per file format.
+
+Every extractor returns the same shape:
+    {"kind": str, "title": str|None, "author": str|None,
+     "extra": dict, "text": str}
+
+kind is one of: document, text, data, other.
+Adding a format means adding one function here and one registry entry —
+nothing else in the system changes.
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+from html.parser import HTMLParser
+from pathlib import Path
+
+MAX_TEXT_CHARS = 200_000
+
+
+def _clean(text: str) -> str:
+    lines = [ln.strip() for ln in text.splitlines()]
+    return "\n".join(ln for ln in lines if ln)[:MAX_TEXT_CHARS]
+
+
+# ---------------------------------------------------------------- text files
+def extract_plain(path: Path) -> dict:
+    raw = path.read_text(encoding="utf-8", errors="replace")
+    title = None
+    for ln in raw.splitlines():
+        if ln.strip():
+            title = ln.strip().lstrip("# ").strip()[:200]
+            break
+    return {"kind": "text", "title": title, "author": None, "extra": {}, "text": _clean(raw)}
+
+
+# ---------------------------------------------------------------- html
+class _HTMLText(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.title: str | None = None
+        self._in_title = False
+        self._skip = 0
+        self.parts: list[str] = []
+
+    def handle_starttag(self, tag, attrs):
+        if tag == "title":
+            self._in_title = True
+        if tag in ("script", "style"):
+            self._skip += 1
+
+    def handle_endtag(self, tag):
+        if tag == "title":
+            self._in_title = False
+        if tag in ("script", "style") and self._skip:
+            self._skip -= 1
+
+    def handle_data(self, data):
+        if self._skip:
+            return
+        if self._in_title:
+            self.title = (self.title or "") + data.strip()
+        text = data.strip()
+        if text:
+            self.parts.append(text)
+
+
+def extract_html(path: Path) -> dict:
+    parser = _HTMLText()
+    parser.feed(path.read_text(encoding="utf-8", errors="replace"))
+    return {
+        "kind": "document",
+        "title": parser.title,
+        "author": None,
+        "extra": {},
+        "text": _clean("\n".join(parser.parts)),
+    }
+
+
+# ---------------------------------------------------------------- docx
+def extract_docx(path: Path) -> dict:
+    import docx  # python-docx
+
+    doc = docx.Document(str(path))
+    text = "\n".join(p.text for p in doc.paragraphs)
+    props = doc.core_properties
+    return {
+        "kind": "document",
+        "title": props.title or None,
+        "author": props.author or None,
+        "extra": {"paragraphs": len(doc.paragraphs)},
+        "text": _clean(text),
+    }
+
+
+# ---------------------------------------------------------------- pdf
+def extract_pdf(path: Path) -> dict:
+    from pypdf import PdfReader
+
+    reader = PdfReader(str(path))
+    text = "\n".join((page.extract_text() or "") for page in reader.pages)
+    meta = reader.metadata or {}
+    return {
+        "kind": "document",
+        "title": (getattr(meta, "title", None) or None),
+        "author": (getattr(meta, "author", None) or None),
+        "extra": {"pages": len(reader.pages)},
+        "text": _clean(text),
+    }
+
+
+# ---------------------------------------------------------------- data files
+_SHARED_CON = None
+
+
+def _shared_con():
+    global _SHARED_CON
+    if _SHARED_CON is None:
+        import duckdb
+
+        _SHARED_CON = duckdb.connect(":memory:")
+    return _SHARED_CON
+
+
+def _duckdb_profile(path: Path, reader_sql: str) -> dict:
+    """Schema + row count + per-column nulls, computed by DuckDB itself.
+
+    Uses one shared in-memory connection across files — opening a fresh
+    DuckDB connection per file costs ~200ms and would dominate scan time.
+    """
+    con = _shared_con()
+    try:
+        cols = con.execute(
+            f"DESCRIBE SELECT * FROM {reader_sql}"
+        ).fetchall()  # (name, type, ...)
+        columns = [{"name": c[0], "type": c[1]} for c in cols]
+        row_count = con.execute(
+            f"SELECT count(*) FROM {reader_sql}"
+        ).fetchone()[0]
+        nulls = {}
+        if columns and row_count:
+            exprs = ", ".join(
+                f'count(*) - count("{c["name"]}")' for c in columns
+            )
+            row = con.execute(
+                f"SELECT {exprs} FROM {reader_sql}"
+            ).fetchone()
+            nulls = {c["name"]: row[i] for i, c in enumerate(columns)}
+        sample_rows = con.execute(
+            f"SELECT * FROM {reader_sql} LIMIT 5"
+        ).fetchall()
+    except Exception:
+        raise
+    # SUMMARIZE: per-column statistics (min/max, approx_unique, avg,
+    # null %) — the profile data-file topics are built from, so sample
+    # values never have to leak into topic computation.
+    summarize = []
+    try:
+        for r in con.execute(f"SUMMARIZE SELECT * FROM {reader_sql}").fetchall():
+            summarize.append({
+                "name": r[0], "type": r[1],
+                "min": None if r[2] is None else str(r[2]),
+                "max": None if r[3] is None else str(r[3]),
+                "approx_unique": r[4],
+                "avg": None if r[5] is None else str(r[5]),
+                "null_percentage": None if r[11] is None else str(r[11]),
+            })
+    except Exception:
+        summarize = []
+    text_parts = [
+        f"columns: {', '.join(c['name'] for c in columns)}",
+        f"rows: {row_count}",
+    ]
+    for r in sample_rows:
+        text_parts.append(" | ".join("" if v is None else str(v) for v in r))
+    return {
+        "kind": "data",
+        "title": None,
+        "author": None,
+        "extra": {
+            "columns": columns,
+            "row_count": row_count,
+            "nulls": nulls,
+            "summarize": summarize,
+        },
+        "text": _clean("\n".join(text_parts)),
+    }
+
+
+def _quoted(path: Path) -> str:
+    return "'" + str(path).replace("'", "''") + "'"
+
+
+def extract_csv(path: Path) -> dict:
+    delim = "\t" if path.suffix.lower() == ".tsv" else ","
+    return _duckdb_profile(path, f"read_csv({_quoted(path)}, header=true, delim='{delim}')")
+
+
+def extract_json_data(path: Path) -> dict:
+    try:
+        return _duckdb_profile(path, f"read_json_auto({_quoted(path)})")
+    except Exception:
+        # Not tabular JSON — treat as text.
+        result = extract_plain(path)
+        result["kind"] = "text"
+        return result
+
+
+def extract_parquet(path: Path) -> dict:
+    return _duckdb_profile(path, f"read_parquet({_quoted(path)})")
+
+
+def extract_xlsx(path: Path) -> dict:
+    import openpyxl
+
+    wb = openpyxl.load_workbook(str(path), read_only=True, data_only=True)
+    sheets = {}
+    text_parts = []
+    total_rows = 0
+    summarize = []
+    for ws in wb.worksheets:
+        rows = list(ws.iter_rows(values_only=True))
+        if not rows:
+            sheets[ws.title] = {"rows": 0, "columns": []}
+            continue
+        header = ["" if h is None else str(h) for h in rows[0]]
+        sheets[ws.title] = {"rows": len(rows) - 1, "columns": header}
+        total_rows += len(rows) - 1
+        text_parts.append(f"sheet: {ws.title}; columns: {', '.join(header)}")
+        for r in rows[1:6]:
+            text_parts.append(" | ".join("" if v is None else str(v) for v in r))
+        if not summarize and rows[1:]:
+            # SUMMARIZE-equivalent for the first sheet, computed from the
+            # rows already in memory: distinct count (capped set), min/max
+            # over comparable values, null percentage.
+            body = rows[1:]
+            for ci, cname in enumerate(header):
+                vals = [r[ci] for r in body if ci < len(r) and r[ci] is not None]
+                distinct = len({str(v) for v in vals[:5000]})
+                nums = [v for v in vals if isinstance(v, (int, float))]
+                pool = nums if nums else [str(v) for v in vals]
+                summarize.append({
+                    "name": cname, "type": "unknown",
+                    "min": None if not pool else str(min(pool)),
+                    "max": None if not pool else str(max(pool)),
+                    "approx_unique": distinct,
+                    "avg": (str(round(sum(nums) / len(nums), 4)) if nums else None),
+                    "null_percentage": str(round(
+                        100.0 * (len(body) - len(vals)) / len(body), 2)),
+                })
+    wb.close()
+    first = next(iter(sheets.values()), {"columns": []})
+    return {
+        "kind": "data",
+        "title": None,
+        "author": None,
+        "extra": {
+            "sheets": sheets,
+            "row_count": total_rows,
+            "columns": [{"name": c, "type": "unknown"} for c in first["columns"]],
+            "nulls": {},
+            "summarize": summarize,
+        },
+        "text": _clean("\n".join(text_parts)),
+    }
+
+
+def extract_sqlite(path: Path) -> dict:
+    con = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
+    try:
+        tables = [
+            r[0]
+            for r in con.execute(
+                "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
+            )
+        ]
+        text_parts = []
+        columns_all = []
+        total = 0
+        summarize = []
+        for t in tables:
+            cols = [r[1] for r in con.execute(f'PRAGMA table_info("{t}")')]
+            n = con.execute(f'SELECT count(*) FROM "{t}"').fetchone()[0]
+            total += n
+            columns_all += cols
+            text_parts.append(f"table: {t}; columns: {', '.join(cols)}; rows: {n}")
+            for c in cols:
+                try:
+                    cnt, uniq, mn, mx, avg = con.execute(
+                        f'SELECT count("{c}"), count(DISTINCT "{c}"), min("{c}"), '
+                        f'max("{c}"), avg("{c}") FROM "{t}"'
+                    ).fetchone()
+                except Exception:
+                    cnt, uniq, mn, mx, avg = con.execute(
+                        f'SELECT count("{c}"), count(DISTINCT "{c}"), min("{c}"), '
+                        f'max("{c}") , NULL FROM "{t}"'
+                    ).fetchone()
+                summarize.append({
+                    "name": c, "type": "unknown", "table": t,
+                    "min": None if mn is None else str(mn),
+                    "max": None if mx is None else str(mx),
+                    "approx_unique": uniq,
+                    "avg": None if avg is None else str(avg),
+                    "null_percentage": (
+                        str(round(100.0 * (n - cnt) / n, 2)) if n else "0"
+                    ),
+                })
+    finally:
+        con.close()
+    return {
+        "kind": "data",
+        "title": None,
+        "author": None,
+        "extra": {
+            "tables": tables,
+            "row_count": total,
+            "columns": [{"name": c, "type": "unknown"} for c in columns_all],
+            "nulls": {},
+            "summarize": summarize,
+        },
+        "text": _clean("\n".join(text_parts)),
+    }
+
+
+def extract_other(path: Path) -> dict:
+    return {"kind": "other", "title": None, "author": None, "extra": {}, "text": ""}
+
+
+EXTRACTORS = {
+    ".txt": extract_plain,
+    ".md": extract_plain,
+    ".html": extract_html,
+    ".htm": extract_html,
+    ".docx": extract_docx,
+    ".pdf": extract_pdf,
+    ".csv": extract_csv,
+    ".tsv": extract_csv,
+    ".json": extract_json_data,
+    ".parquet": extract_parquet,
+    ".xlsx": extract_xlsx,
+    ".sqlite": extract_sqlite,
+    ".db": extract_sqlite,
+}
+
+SUPPORTED_EXTS = set(EXTRACTORS)
+
+
+def extract(path: Path) -> dict:
+    fn = EXTRACTORS.get(path.suffix.lower(), extract_other)
+    try:
+        return fn(path)
+    except Exception as exc:  # a broken file must never kill a scan
+        return {
+            "kind": "other",
+            "title": None,
+            "author": None,
+            "extra": {"extract_error": str(exc)[:300]},
+            "text": "",
+        }

folderlake/hierarchy.py

@@ -0,0 +1,297 @@
+"""Hierarchical topic tree — agglomerative clustering per file kind.
+
+Flat per-file tags stay as metadata; the *organisation* of topics is
+this tree. Files of each kind (document, data, text) are clustered
+separately — the kinds are never mixed, per the product decision that
+their browsing scenarios differ.
+
+Method: corpus TF-IDF vectors from the stored per-file term counts
+(log TF x IDF, L2-normalised, light DF hygiene), then agglomerative
+clustering with centroid linkage. The merge tree becomes a display
+tree of at most MAX_DEPTH levels: a sub-cluster earns its own node
+when it holds at least MIN_GROUP files (smaller branches simply stay
+covered by their parent's member list), and a single child that holds
+essentially all of its parent's files replaces the parent instead of
+adding an empty level. Node labels are the cluster's top phrase terms
+(the same bonus/suppression rules as per-file tags, in topics.py).
+
+topic_node_files stores each node's FULL subtree membership, so
+filtering by any node is a single lookup.
+
+Scale honesty: centroid agglomeration is O(N^2) pairwise work, so a
+kind with more than MAX_FILES_PER_KIND files gets no tree (the viewer
+falls back to the flat topic list) rather than a slow scan.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+from pathlib import Path
+
+import duckdb
+
+from .topics import _select_tags, _term_bonus
+
+MIN_GROUP = 3
+MAX_DEPTH = 4
+# A cluster bigger than this whose halves both qualify as groups is
+# split at the SAME display level instead of becoming one wide node —
+# otherwise a binary merge tree gives every level at most 2 children
+# and the tree stays sparse.
+WIDE_SPLIT = 12
+MAX_FILES_PER_KIND = 400
+
+KIND_LABELS = {"document": "Documents", "data": "Data files", "text": "Texts"}
+
+
+# ----------------------------------------------------------------------
+# Vectors
+# ----------------------------------------------------------------------
+def _vectors(terms_by_file: dict[str, Counter]) -> dict[str, dict[str, float]]:
+    n = len(terms_by_file)
+    df: Counter = Counter()
+    for counts in terms_by_file.values():
+        df.update(counts.keys())
+    vectors: dict[str, dict[str, float]] = {}
+    for fid, counts in terms_by_file.items():
+        vec: dict[str, float] = {}
+        for term, cnt in counts.items():
+            d = df[term]
+            if n >= 10 and (d < 2 or d > 0.7 * n):
+                continue  # DF hygiene: hapax noise + corpus wallpaper
+            idf = math.log((n + 1) / (d + 1)) + 1.0
+            vec[term] = math.log1p(cnt) * idf
+        norm = math.sqrt(sum(v * v for v in vec.values())) or 1.0
+        vectors[fid] = {t: v / norm for t, v in vec.items()}
+    return vectors
+
+
+def _cosine(a: dict[str, float], b: dict[str, float]) -> float:
+    if len(a) > len(b):
+        a, b = b, a
+    return sum(v * b.get(t, 0.0) for t, v in a.items())
+
+
+# ----------------------------------------------------------------------
+# Agglomerative clustering (centroid linkage)
+# ----------------------------------------------------------------------
+class _Cluster:
+    __slots__ = ("members", "total", "left", "right", "sim")
+
+    def __init__(self, members, total, left=None, right=None, sim=0.0):
+        self.members = members          # file ids in this subtree
+        self.total = total              # summed member vectors (dict)
+        self.left, self.right, self.sim = left, right, sim
+
+    def centroid(self) -> dict[str, float]:
+        norm = math.sqrt(sum(v * v for v in self.total.values())) or 1.0
+        inv = 1.0 / norm
+        return {t: v * inv for t, v in self.total.items()}
+
+
+def _agglomerate(vectors: dict[str, dict[str, float]]) -> _Cluster:
+    active: dict[int, _Cluster] = {
+        i: _Cluster([fid], dict(vec))
+        for i, (fid, vec) in enumerate(vectors.items())
+    }
+    sims: dict[tuple[int, int], float] = {}
+
+    def sim(i: int, j: int) -> float:
+        key = (i, j) if i < j else (j, i)
+        if key not in sims:
+            sims[key] = _cosine(active[i].centroid(), active[j].centroid())
+        return sims[key]
+
+    next_id = len(active)
+    while len(active) > 1:
+        ids = list(active)
+        best, best_pair = -2.0, None
+        for x in range(len(ids)):
+            for y in range(x + 1, len(ids)):
+                s = sim(ids[x], ids[y])
+                if s > best:
+                    best, best_pair = s, (ids[x], ids[y])
+        i, j = best_pair
+        a, b = active.pop(i), active.pop(j)
+        total = dict(a.total)
+        for t, v in b.total.items():
+            total[t] = total.get(t, 0.0) + v
+        active[next_id] = _Cluster(a.members + b.members, total,
+                                   left=a, right=b, sim=max(best, 0.0))
+        sims = {k: v for k, v in sims.items() if i not in k and j not in k}
+        next_id += 1
+    return next(iter(active.values()))
+
+
+# ----------------------------------------------------------------------
+# Display tree
+# ----------------------------------------------------------------------
+def _label(cluster: _Cluster, corpus_total: dict[str, float],
+           corpus_size: float) -> str:
+    """Label by lift: terms over-represented in this cluster versus the
+    whole kind corpus. Raw frequency crowns corpus wallpaper ('more',
+    header boilerplate) that appears everywhere at the same rate; lift
+    sends those to zero and surfaces what the group is actually about.
+    """
+    cluster_size = sum(cluster.total.values()) or 1.0
+    ranked = []
+    for term, value in cluster.total.items():
+        share_in = value / cluster_size
+        share_all = corpus_total.get(term, 0.0) / (corpus_size or 1.0)
+        lift = math.log((share_in + 1e-9) / (share_all + 1e-9))
+        if lift > 0:
+            ranked.append((term, value * lift * _term_bonus(term)))
+    ranked.sort(key=lambda kv: -kv[1])
+    top = _select_tags(ranked, Counter(cluster.total), limit=3)
+    # Display polish: a bare "sex" (the dataset column) reads as a
+    # content warning in a topic label. If the other picked terms or
+    # the runners-up can carry the label, drop it; keep it only when
+    # nothing else is available.
+    picked = [t for t, _ in top]
+    if "sex" in picked and len(picked) > 1:
+        picked = [t for t in picked if t != "sex"]
+        for t, _ in ranked:
+            if len(picked) >= 3:
+                break
+            if t not in picked and t != "sex":
+                picked.append(t)
+    return " · ".join(picked) or "misc"
+
+
+def _should_flatten(cluster: _Cluster) -> bool:
+    """A forced marriage: two internally-cohesive halves that barely
+    resemble each other. Such a node gets no display level — its
+    children are promoted to its parent instead. (A binary dendrogram
+    otherwise invents meaningless top groups, e.g. finance+garden
+    labelled 'finance'.)"""
+    if cluster.left is None or cluster.right is None:
+        return False
+    internal = min(cluster.left.sim, cluster.right.sim)
+    return internal > 0 and cluster.sim < 0.4 * internal
+
+
+def _child_clusters(cluster: _Cluster, depth: int) -> list[_Cluster]:
+    """The frontier of sub-clusters that become nodes at depth + 1.
+
+    Walk down the merge tree from this cluster's halves: a branch
+    becomes a node when it is a qualifying group; we keep descending
+    through branches that are forced marriages (see _should_flatten)
+    or simply too wide to be one topic (WIDE_SPLIT).
+    """
+    if cluster.left is None or depth >= MAX_DEPTH:
+        return []
+    out: list[_Cluster] = []
+    stack = [cluster.left, cluster.right]
+    while stack:
+        c = stack.pop()
+        if len(c.members) < MIN_GROUP:
+            continue
+        if c.left is not None and (
+                _should_flatten(c) or len(c.members) > WIDE_SPLIT):
+            stack.extend([c.left, c.right])
+        else:
+            out.append(c)
+    return out
+
+
+def build_tree_for_kind(root_cluster: _Cluster, kind: str,
+                        corpus_total: dict[str, float] | None = None
+                        ) -> list[dict]:
+    """Flatten a merge tree into display nodes.
+
+    Each node dict: prefixed id, parent id (or None), kind, label,
+    depth, members (FULL subtree file ids), cohesion.
+    """
+    if corpus_total is None:
+        corpus_total = dict(root_cluster.total)
+    corpus_size = sum(corpus_total.values()) or 1.0
+    nodes: list[dict] = []
+
+    def emit(cluster: _Cluster, parent_id: str | None, depth: int,
+             label: str, cohesion: float) -> None:
+        kids = _child_clusters(cluster, depth)
+        # Chain collapse: one child holding (nearly) every file means
+        # this level adds nothing — emit the child in our place. The
+        # kind root is exempt: it carries the kind label.
+        if parent_id is not None and len(kids) == 1 and \
+                len(kids[0].members) >= len(cluster.members) - (MIN_GROUP - 1):
+            emit(kids[0], parent_id, depth,
+                 _label(kids[0], corpus_total, corpus_size), kids[0].sim)
+            return
+        node_id = f"{kind}-{len(nodes) + 1}"
+        nodes.append({
+            "id": node_id, "parent": parent_id, "kind": kind,
+            "label": label, "depth": depth,
+            "members": list(cluster.members),
+            "cohesion": round(cohesion, 3),
+        })
+        for c in kids:
+            emit(c, node_id, depth + 1,
+                 _label(c, corpus_total, corpus_size), c.sim)
+
+    emit(root_cluster, None, 0, KIND_LABELS.get(kind, kind), 1.0)
+    return nodes
+
+
+# ----------------------------------------------------------------------
+# Build + store
+# ----------------------------------------------------------------------
+def build_hierarchy(db_path: Path) -> dict:
+    con = duckdb.connect(str(db_path))
+    try:
+        con.execute(
+            "CREATE TABLE IF NOT EXISTS topic_nodes ("
+            "id VARCHAR PRIMARY KEY, parent_id VARCHAR, kind VARCHAR, "
+            "label VARCHAR, depth INTEGER, file_count INTEGER, cohesion DOUBLE)"
+        )
+        con.execute(
+            "CREATE TABLE IF NOT EXISTS topic_node_files ("
+            "node_id VARCHAR, file_id VARCHAR)"
+        )
+        terms: dict[str, dict[str, Counter]] = {}
+        for fid, kind, term, cnt in con.execute(
+            "SELECT f.id, f.kind, t.term, t.cnt FROM files f "
+            "JOIN file_terms t ON t.file_id = f.id"
+        ).fetchall():
+            terms.setdefault(kind, {}).setdefault(fid, Counter())[term] = cnt
+        all_files: dict[str, list[str]] = {}
+        for fid, kind in con.execute("SELECT id, kind FROM files").fetchall():
+            all_files.setdefault(kind, []).append(fid)
+
+        node_rows, member_rows = [], []
+        skipped: dict[str, int] = {}
+        for kind, file_ids in sorted(all_files.items()):
+            if len(file_ids) > MAX_FILES_PER_KIND:
+                skipped[kind] = len(file_ids)
+                continue
+            terms_by_file = terms.get(kind, {})
+            vectors = _vectors({fid: terms_by_file.get(fid, Counter())
+                                for fid in file_ids})
+            corpus_total: dict[str, float] = {}
+            for vec in vectors.values():
+                for t, v in vec.items():
+                    corpus_total[t] = corpus_total.get(t, 0.0) + v
+            tree = build_tree_for_kind(_agglomerate(vectors), kind,
+                                       corpus_total)
+            for node in tree:
+                node_rows.append((
+                    node["id"], node["parent"], kind, node["label"],
+                    node["depth"], len(node["members"]), node["cohesion"],
+                ))
+                member_rows.extend(
+                    (node["id"], fid) for fid in node["members"])
+        con.execute("BEGIN TRANSACTION")
+        con.execute("DELETE FROM topic_nodes")
+        con.execute("DELETE FROM topic_node_files")
+        if node_rows:
+            con.executemany("INSERT INTO topic_nodes VALUES (?,?,?,?,?,?,?)",
+                            node_rows)
+        if member_rows:
+            con.executemany("INSERT INTO topic_node_files VALUES (?,?)",
+                            member_rows)
+        con.execute("COMMIT")
+        return {"nodes": len(node_rows), "kinds": sorted(all_files),
+                "skipped_too_large": skipped}
+    finally:
+        con.close()