Topic: Tool to convert parent chains to pools

import json
import os.path
import requests as requests
from ratelimit import limits, sleep_and_retry


@sleep_and_retry
@limits(calls=1, period=1)
def api_limit():
    pass


def get_login() -> tuple[str, str]:
    if os.path.isfile("login.json"):
        with open("login.json", "r") as file:
            data = json.load(file)
            return data["username"], data["key"]
    username = input("Username: ")
    key = input("API key: ")
    if input("Save login (y/n): ") == "y":
        with open("login.json", "w") as file:
            json.dump({"username": username, "key": key}, file)
    return username, key


def get_page(session: requests.Session, post_id: str) -> dict:
    api_limit()
    response = session.get(f"https://e621.net/posts/{post_id}.json")
    if not response:
        print(f"Unable to get {post_id}: Status {response.status_code}, response {response.text}")
        exit(1)
    post = response.json()["post"]
    return {"id": post_id, "parent": post["relationships"]["parent_id"], "pools": post["pools"]}


def get_chain(session: requests.Session, start: str) -> list:
    print("Getting posts")
    post_id = start
    result = []
    while post_id is not None:
        result.append(get_page(session, post_id))
        print(f"Got {post_id}, pools: {result[-1]['pools']}")
        post_id = result[-1]["parent"]
    result.reverse()
    return result


def get_category():
    while True:
        response = input("category (series/collection): ")
        if response == "series" or response == "s":
            return "series"
        if response == "collection" or response == "c":
            return "collection"
        print("Invalid response")


def create_pool(session: requests.Session, chain: list):
    payload = dict()
    payload["pool[name]"] = input("name: ")
    payload["pool[description]"] = input("description: ")
    payload["pool[category]"] = get_category()
    payload["pool[is_active]"] = "true"
    payload["pool[post_ids]"] = " ".join([str(post["id"]) for post in chain])
    api_limit()
    response = session.post("https://e621.net/pools.json", data=payload)
    if response:
        print(f"Created pool: https://e621.net/pools/{response.json()['id']}")
    else:
        print(f"Failed to create pool: Status {response.status_code}, response {response.text}")


def remove_parents(session: requests.Session, chain: list):
    for post in chain[1:]:
        payload = dict()
        payload["post[parent_id]"] = ""
        payload["post[old_parent_id]"] = str(post["parent"])
        payload["post[edit_reason]"] = "Convert chain to pool (automatic update)"
        api_limit()
        response = session.patch(f"https://e621.net/posts/{post['id']}.json", data=payload)
        if response:
            print(f"Removed parent from post {post['id']}")
        else:
            print(f"Failed to remove parent from post {post['id']}: Status {response.status_code}, response {response.text}")


def main():
    login = get_login()
    start = input("ID of last post in chain: ")
    with requests.Session() as session:
        session.auth = login
        session.headers.update({"User-Agent": "ChainToPool/1.0 (by SCTH)"})
        chain = get_chain(session, start)
        if input("Create pool (y/n): ") == "y":
            create_pool(session, chain)
        if input("Remove parents from chain (y/n): ") == "y":
            remove_parents(session, chain)


if __name__ == '__main__':
    main()

804034
└── 805199
    └── 810075
        ├── 820373
        ├── 823457
        │   └── 826198
        │       ├── 836150
        │       │   └── 837626
        │       │       └── 841256
        │       │           └── 843983
        │       │               └── 848483
        │       │                   └── 859762
        │       │                       └── 861821
        │       │                           └── 863266
        │       │                               └── 864734
        │       └── 880351
        │           └── 885483
        │               ├── 895391
        │               └── 897532
        │                   ├── 901080
        │                   └── 905829
        │                       ├── 906744
        │                       │   └── 917588
        │                       │       └── 922286
        │                       └── 910073
        │                           └── 915656
        │                               └── 937622
        │                                   └── 947714
        │                                       └── 980574
        │                                           └── 1221162
        │                                               └── 1223126
        │                                                   └── 1286083
        └── 843166
            └── 894438
                └── 898088

import threading
import requests
import os
import re
import polars
import zlib

### Export downloader
print_lock = threading.Lock()
headers = {'User-Agent': "ExportDownloader/0.3 (by SCTH)"}

def locked_print(*args, **kwargs):
    with print_lock:
        print(*args, **kwargs)


def download_and_extract(url, filename):
    locked_print(f"Downloading {url}")
    decompressor = zlib.decompressobj(wbits=zlib.MAX_WBITS | 16)
    with requests.get(url, headers=headers, stream=True) as resp, open(filename, "wb") as f_out:
        for chunk in resp.iter_content(chunk_size=8192):
            chunk_inflated = decompressor.decompress(chunk)
            f_out.write(chunk_inflated)
    locked_print(f"Download of {url} finished")


def find_best_csv(name):
    files = [f for f in os.listdir(".") if os.path.isfile(f) and f.startswith(name) and f.endswith(".csv")]
    if len(files) == 0:
        return None
    return max(files)


def get_exports(export_names):
    export_page = requests.get("https://e621.net/db_export/", headers=headers).text
    current_bests = [find_best_csv(name) for name in export_names]
    new_bests = [max(re.findall(f"{name}-....-..-..\\.csv", export_page)) for name in export_names]
    threads = []
    for current, new, name in zip(current_bests, new_bests, export_names):
        if current != new:
            if current is not None:
                os.remove(current)
            threads.append(threading.Thread(target=download_and_extract,
                                            args=(f"https://e621.net/db_export/{new}.gz", new)))
            threads[-1].start()
    for thread in threads:
        thread.join()
    return new_bests


### Main program - find and print post trees
allow_deleted = False # Include deleted posts in the search
allow_in_pool = False # Include posts that are already in a pool
allow_non_chain = False # Include cases where a single post has multiple children


def print_tree(posts, pools, relations, root, depth=0):
    print((" " * depth) + str(root), end='')
    if not posts.filter(polars.col("id")==root, polars.col("is_deleted")=="t").is_empty():
        print(" (deleted)", end='')
    for row in pools.filter(polars.col("post_id")==root).iter_rows(named=True):
        print(f" ({row["pool_id"]})", end='')
    print()
    for row in relations.filter(direct_parent=root).iter_rows(named=True):
        print_tree(posts, pools, relations, row["child_id"], depth + 1)


def main():
    posts_file, pools_file = get_exports(["posts", "pools"])
    print("Reading pools csv")
    pools = polars.read_csv(pools_file).select("id",
                                               post_ids=polars.col("post_ids").str.strip_chars("{}").str.split(",")) \
        .explode("post_ids") \
        .filter(polars.col("post_ids") != "") \
        .select(pool_id="id", post_id=polars.col("post_ids").str.to_integer())
    print("Reading posts csv")
    posts = polars.read_csv(posts_file).select("id", "parent_id", "is_deleted")
    if not allow_deleted:
        posts = posts.filter(is_deleted="f")
    if not allow_in_pool:
        posts = posts.join(other=pools, how="anti", left_on="id", right_on="post_id")
    print("Read, building list")
    roots = posts.filter(polars.col("parent_id").is_null())
    children = posts.filter(polars.col("parent_id").is_not_null())
    level_frame = roots.join(other=children, how="inner", left_on="id", right_on="parent_id", suffix="_child")\
        .select(root_id="id", direct_parent="id", child_id="id_child")
    if not allow_non_chain:
        level_frame = level_frame.unique(subset="direct_parent", keep="none")
    level = 1
    relations = level_frame.clear()
    while level_frame.shape[0] > 0:
        relations.vstack(level_frame, in_place=True)
        level += 1
        level_frame = level_frame.join(other=children, how="inner", left_on="child_id", right_on="parent_id")\
            .select(root_id="root_id", direct_parent="child_id", child_id="id")
        if not allow_non_chain:
            level_frame = level_frame.unique(subset="direct_parent", keep="none")
    relations.rechunk()
    print("Calculating tree sizes")
    tree_sizes = relations.group_by("root_id").len().sort("len", descending=True)
    for row in tree_sizes.iter_rows(named=True):
        show = input(f"Tree of size {row["len"] + 1} at {row["root_id"]}: ")
        if len(show) > 0 and show[0] == "y":
            print_tree(posts, pools, relations, row["root_id"])


if __name__ == '__main__':
    main()