357 lines
12 KiB
Python
357 lines
12 KiB
Python
from datetime import datetime as dt
|
|
import json
|
|
import sys
|
|
from time import sleep
|
|
import traceback
|
|
from typing import Any, Dict, List
|
|
|
|
import pytumblr
|
|
import polars as pl
|
|
|
|
|
|
def filter_for_after(df: pl.DataFrame,
|
|
after: dt) -> pl.DataFrame:
|
|
if df.is_empty() or not after: return pl.DataFrame()
|
|
base_col = pl.col('date')
|
|
if df.get_column('date').dtype == pl.String:
|
|
base_col = pl.when(base_col.str.contains('GMT')) \
|
|
.then(base_col.str.to_date('%Y-%m-%d %H:%M:%S GMT')) \
|
|
.otherwise(base_col.str.to_date('%Y-%m-%d %H:%M:%S'))
|
|
return df.filter(base_col >= after.date())
|
|
|
|
|
|
def filter_for_before(df: pl.DataFrame,
|
|
before: dt) -> pl.DataFrame:
|
|
if df.is_empty() or not before: return pl.DataFrame()
|
|
base_col = pl.col('date')
|
|
if df.get_column('date').dtype == pl.String:
|
|
base_col = pl.when(base_col.str.contains('GMT')) \
|
|
.then(base_col.str.to_date('%Y-%m-%d %H:%M:%S GMT')) \
|
|
.otherwise(base_col.str.to_date('%Y-%m-%d %H:%M:%S'))
|
|
return df.filter(base_col <= before.date())
|
|
|
|
|
|
def handle_error(e: Exception, data: Any, df: pl.DataFrame) -> None:
|
|
print(e)
|
|
traceback.print_exc()
|
|
schema_dict = {col: str(dtype) for col, dtype in df.schema.items()}
|
|
with open("schema.json", "w") as f:
|
|
json.dump(schema_dict, f, indent=4)
|
|
with open("bad_posts.json", "w") as f:
|
|
json.dump(data, f, indent=4, default=str)
|
|
|
|
|
|
def handle_limit_exceeded(data) -> bool:
|
|
if 'meta' in data and data['meta']['status'] == 429:
|
|
print("WARN: API limit exceeded. Sleeping for 10 minutes...")
|
|
sleep(600)
|
|
return True
|
|
sleep(3)
|
|
return False
|
|
|
|
|
|
def fetch_minimized_post(blog_name: str,
|
|
client: pytumblr.TumblrRestClient,
|
|
id_string: str) -> Dict[str, Any]:
|
|
api_url = f"/v2/blog/{blog_name}/posts"
|
|
data: Dict[str, Any] = client.send_api_request(method='get',
|
|
url=api_url,
|
|
params={'id': id_string},
|
|
valid_parameters=['id'])
|
|
|
|
if 'posts' in data:
|
|
posts = data['posts']
|
|
if handle_limit_exceeded(data):
|
|
data = client.send_api_request(method='get', url=api_url, params={'id': id_string}, valid_parameters=['id'])
|
|
posts = data['posts']
|
|
if 'posts' not in data or not data['posts']:
|
|
print(f"WARN: Post with ID {id_string} could not be retrieved!")
|
|
print(data)
|
|
return {}
|
|
|
|
minimized_post: Dict[str, Any] = {}
|
|
if posts:
|
|
minimized_post = posts[0]
|
|
minimized_post = {
|
|
'id_string': minimized_post['id_string'],
|
|
'url': minimized_post['post_url'],
|
|
'date_str': minimized_post['date'],
|
|
'tags': minimized_post['tags'],
|
|
'is_og': 'parent_post_url' not in minimized_post
|
|
}
|
|
return minimized_post
|
|
|
|
|
|
def fetch_queued_posts(client: pytumblr.TumblrRestClient,
|
|
args: Dict[str, Any]) -> pl.DataFrame:
|
|
queued_posts: pl.DataFrame = pl.DataFrame()
|
|
blog_name: str = args['blog']
|
|
|
|
offset: int = 0
|
|
limit: int = 20
|
|
|
|
while True:
|
|
# Begin LOOP
|
|
# Get me some queued posts via REST! 😈🍪🍪🍪
|
|
data: Dict[str, Any]
|
|
data = client.queue(f"{blog_name}.tumblr.com",
|
|
offset=offset,
|
|
limit=limit)
|
|
|
|
if handle_limit_exceeded(data):
|
|
data = client.queue(f"{blog_name}.tumblr.com",
|
|
offset=offset,
|
|
limit=limit)
|
|
|
|
# This struct breaks Polars for some reason 🤷
|
|
posts_json = data['posts']
|
|
for post in posts_json:
|
|
if 'trail' in post:
|
|
del post['trail']
|
|
|
|
try:
|
|
curr_posts: pl.DataFrame = pl.DataFrame(posts_json)
|
|
except Exception as e:
|
|
handle_error(e, data, queued_posts)
|
|
sys.exit(0)
|
|
|
|
# Stop the presses if we found no posts.
|
|
if curr_posts.is_empty() or len(curr_posts) < 1:
|
|
print('Stopping queued post retrieval, as no more queued posts were found.')
|
|
break
|
|
|
|
queued_posts = pl.concat([queued_posts, curr_posts], how="diagonal_relaxed")
|
|
|
|
offset += limit
|
|
# End LOOP
|
|
|
|
# Empty check.
|
|
if queued_posts.is_empty():
|
|
print("No queued posts were found.")
|
|
|
|
return queued_posts
|
|
|
|
|
|
def fetch_drafts(client: pytumblr.TumblrRestClient,
|
|
args: Dict[str, Any]) -> pl.DataFrame:
|
|
drafts: pl.DataFrame = pl.DataFrame()
|
|
blog_name: str = args['blog']
|
|
draft_url = f"/v2/blog/{blog_name}/posts/draft"
|
|
|
|
while True:
|
|
# Begin LOOP
|
|
# Get me some drafts via REST! 😈🍪🍪🍪
|
|
data: Dict[str, Any]
|
|
data = client.send_api_request("get", draft_url)
|
|
|
|
if handle_limit_exceeded(data):
|
|
data = client.send_api_request("get", draft_url)
|
|
|
|
# This struct breaks Polars for some reason 🤷
|
|
drafts_json = data['posts']
|
|
for draft in drafts_json:
|
|
if 'trail' in draft:
|
|
del draft['trail']
|
|
|
|
curr_drafts: pl.DataFrame = pl.DataFrame(drafts_json)
|
|
drafts = pl.concat([drafts, curr_drafts], how="diagonal_relaxed")
|
|
|
|
if '_links' in data:
|
|
draft_url = data['_links']['next']['href']
|
|
continue
|
|
else:
|
|
print('Stopping draft retrieval, as no more drafts were found.')
|
|
break
|
|
# End LOOP
|
|
|
|
# Empty check.
|
|
if drafts.is_empty():
|
|
print("No drafts were found.")
|
|
|
|
return drafts
|
|
|
|
|
|
def fetch_posts(client: pytumblr.TumblrRestClient,
|
|
args: Dict[str, Any]) -> pl.DataFrame:
|
|
posts: pl.DataFrame = pl.DataFrame()
|
|
blog_name: str = args['blog']
|
|
|
|
# We populate params, starting with any tags for filtering.
|
|
params = {}
|
|
if args['tags']:
|
|
params.update({'tag': ','.join(args['tags'])})
|
|
|
|
total: int = 0
|
|
offset: int = 0
|
|
limit: int = 20
|
|
|
|
# The request loop that pulls all data from the APIs.
|
|
while True:
|
|
# Begin LOOP
|
|
# Get me some posts via REST! 😈🍪🍪🍪
|
|
data: Dict[str, Any]
|
|
data = client.posts(f"{blog_name}.tumblr.com",
|
|
offset=offset,
|
|
limit=limit,
|
|
**params)
|
|
|
|
if handle_limit_exceeded(data):
|
|
data = client.posts(f"{blog_name}.tumblr.com",
|
|
offset=offset,
|
|
limit=limit,
|
|
**params)
|
|
|
|
# These structs break Polars for some reason 🤷
|
|
posts_json = data['posts']
|
|
for post in posts_json:
|
|
if 'trail' in post:
|
|
del post['trail']
|
|
if 'player' in post:
|
|
del post['player']
|
|
|
|
try:
|
|
curr_posts: pl.DataFrame = pl.DataFrame(posts_json, strict=False, infer_schema_length=None)
|
|
except Exception as e:
|
|
handle_error(e, data, curr_posts)
|
|
sys.exit(1)
|
|
|
|
# Stop the presses if we found no posts.
|
|
if curr_posts.is_empty() or len(curr_posts) < 1:
|
|
print('Stopping, as no more posts were found.')
|
|
break
|
|
|
|
next_off: int = 0
|
|
|
|
# Total init check for the first iteration, but always checked for sanity.
|
|
if not total and 'total_posts' in data:
|
|
total_posts = data['total_posts']
|
|
print(f"I'm working with {total_posts} total posts...")
|
|
total = total_posts
|
|
|
|
if args['since'] and not args['until']:
|
|
after: dt = args['since']
|
|
curr_posts = filter_for_after(curr_posts, after)
|
|
if curr_posts.is_empty():
|
|
print(
|
|
f"All posts since {after.year}-{after.month}-{after.day} retrieved.")
|
|
return posts
|
|
if args['until'] and not args['since']:
|
|
before: dt = args['until']
|
|
curr_posts = filter_for_before(curr_posts, before)
|
|
if curr_posts.is_empty():
|
|
print(
|
|
f"All posts until {before.year}-{before.month}-{before.day} retrieved.")
|
|
return posts
|
|
if args['since'] and args['until']:
|
|
if not posts.is_empty() and curr_posts.is_empty():
|
|
print(f"All posts since {after.year}-{after.month}-{after.day}, "
|
|
+ f"until {before.year}-{before.month}-{before.day} retrieved.")
|
|
return posts
|
|
|
|
# Update the frame with what we found.
|
|
try:
|
|
posts = pl.concat([posts, curr_posts], how="diagonal_relaxed")
|
|
except Exception as e:
|
|
handle_error(e, data, posts)
|
|
sys.exit(1)
|
|
|
|
# The increment and status printing.
|
|
if next_off != 0 and next_off != offset:
|
|
offset = next_off
|
|
else:
|
|
offset += limit
|
|
if not (args['since'] or args['until']) and total:
|
|
print(
|
|
f"Processed batch {offset // limit} of {(total // 20) + 1}...")
|
|
# End LOOP
|
|
|
|
if posts.is_empty():
|
|
print("No posts were found.")
|
|
|
|
# Return posts (both original and not original).
|
|
return posts.unique(subset=["id_string"])
|
|
|
|
|
|
def fetch_activity(client: pytumblr.TumblrRestClient,
|
|
args: Dict[str, Any],
|
|
posts: pl.DataFrame = pl.DataFrame()) -> pl.DataFrame:
|
|
activity: pl.DataFrame = pl.DataFrame(infer_schema_length=None)
|
|
blog_name: str = args['blog']
|
|
|
|
before: dt = args['until']
|
|
after: dt = args['since']
|
|
|
|
activity_url = f"/v2/blog/{blog_name}/notifications"
|
|
if before:
|
|
activity_url += f"?before={str(int(before.timestamp()))}"
|
|
|
|
params = {}
|
|
params.update({'rollups': 'false'})
|
|
|
|
post_map: Dict[str, Any] = {}
|
|
if not posts.is_empty():
|
|
post_map = { row["id"]: row for row in posts.to_dicts() }
|
|
|
|
data: Dict[str, Any]
|
|
while True:
|
|
# Begin LOOP
|
|
# Get me some activity via REST! 😈🍪🍪🍪
|
|
data = client.send_api_request("get", activity_url)
|
|
if handle_limit_exceeded(data):
|
|
data = client.send_api_request("get", activity_url)
|
|
|
|
try:
|
|
curr_activity: pl.DataFrame = pl.DataFrame(data['notifications'], infer_schema_length=None)
|
|
except Exception as e:
|
|
handle_error(e, data, activity)
|
|
sys.exit(0)
|
|
curr_activity = curr_activity.with_columns(pl.from_epoch('timestamp').alias('date'))
|
|
if before:
|
|
curr_activity = filter_for_before(curr_activity, before)
|
|
if after:
|
|
curr_activity = filter_for_after(curr_activity, after)
|
|
|
|
# Add additional post info to notifications.
|
|
# Yeah, df to dictionary list. Sorry. :c
|
|
curr_activity_dicts: List[Dict[str, Any]] = curr_activity.to_dicts()
|
|
for notification in curr_activity_dicts:
|
|
if 'target_post_id' in notification and notification['target_post_id']:
|
|
target_post_id = notification['target_post_id']
|
|
post: Any
|
|
if target_post_id in post_map:
|
|
post = post_map.get(target_post_id)
|
|
notification['post'] = post
|
|
continue
|
|
elif 'type' in notification and notification['type'] == 'post_attribution':
|
|
notification['post'] = {}
|
|
continue
|
|
else:
|
|
post = fetch_minimized_post(blog_name, client, notification['target_post_id'])
|
|
notification['post'] = post
|
|
post_map.update({target_post_id: post})
|
|
continue
|
|
|
|
curr_activity = pl.DataFrame(curr_activity_dicts, strict=False, infer_schema_length=None)
|
|
activity = pl.concat([activity, curr_activity], how="diagonal_relaxed")
|
|
|
|
if '_links' in data:
|
|
next_time = dt.fromtimestamp(int(data['_links']['next']['query_params']['before']))
|
|
if before or after:
|
|
print(f"Reached activity time: {next_time}...")
|
|
if (before and next_time >= before) or (after and next_time <= after):
|
|
print('Stopping activity retrieval, as no more activity is left in date range.')
|
|
break
|
|
activity_url = data['_links']['next']['href']
|
|
continue
|
|
else:
|
|
print('Stopping activity retrieval, as no more activity was found.')
|
|
break
|
|
# End LOOP
|
|
|
|
# Empty check.
|
|
if activity.is_empty():
|
|
print("No activity was found.")
|
|
else:
|
|
activity = activity.unique(subset=["id"])
|
|
|
|
return activity |