149 lines
6.1 KiB
Python
149 lines
6.1 KiB
Python
import argparse
|
|
from dataclasses import asdict
|
|
import dataclasses
|
|
from datetime import datetime as dt
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
from typing import Any, Dict, List
|
|
|
|
import pytumblr
|
|
import polars as pl
|
|
import fetcher
|
|
from model import StatsModel
|
|
|
|
|
|
def get_args() -> Dict[str, Any]:
|
|
"""Pull arguments from command line, turn them into a dictionary of <arg, value>"""
|
|
parser: argparse.ArgumentParser = argparse.ArgumentParser(
|
|
prog='tumblr_stats.py',
|
|
description='Use pytumblr and Polars to calculate stats after setting these enviroment variables: '
|
|
+ '$TUMBLR_CONSUMER_KEY, $TUMBLR_CONSUMER_SECRET, $TUMBLR_OAUTH_TOKEN, and $TUMBLR_OAUTH_SECRET',
|
|
epilog='— Be gay and do crime')
|
|
parser.add_argument('-b', '--blog', type=str, required=True,
|
|
help='blog name for which to calculate stats')
|
|
parser.add_argument('-t', '--tags', type=str, nargs='+',
|
|
help='tag(s) on which to filter (only affects posts)')
|
|
parser.add_argument('-s', '--suppress', type=str, nargs='+',
|
|
choices=['none', 'posts', 'queued', 'drafts', 'activity'],
|
|
default='activity',
|
|
help='objects to skip when aggregating data; will render some stats empty')
|
|
parser.add_argument('--until', type=lambda s: dt.strptime(s, '%Y-%m-%d'),
|
|
help='only gather posts and activity until YYYY-MM-DD')
|
|
parser.add_argument('--since', type=lambda s: dt.strptime(s, '%Y-%m-%d'),
|
|
help='only fetch posts and activity since YYYY-MM-DD')
|
|
parser.add_argument('--queued-file', type=str,
|
|
help='read queued posts from this path instead of fetching them from APIs')
|
|
parser.add_argument('--drafts-file', type=str,
|
|
help='read drafts from this path instead of fetching them from APIs')
|
|
parser.add_argument('--posts-file', type=str,
|
|
help='read posts from this path instead of fetching them from APIs')
|
|
parser.add_argument('--activity-file', type=str,
|
|
help='read activity from this path instead of fetching it from APIs')
|
|
return vars(parser.parse_args())
|
|
|
|
|
|
def init_client() -> pytumblr.TumblrRestClient:
|
|
consumer_key = os.getenv('TUMBLR_CONSUMER_KEY')
|
|
consumer_secret = os.getenv('TUMBLR_CONSUMER_SECRET')
|
|
oauth_token = os.getenv('TUMBLR_OAUTH_TOKEN')
|
|
oauth_secret = os.getenv('TUMBLR_OAUTH_SECRET')
|
|
|
|
missing_vars: List[str] = [name for name,
|
|
val in [('$TUMBLR_CONSUMER_KEY', consumer_key),
|
|
('$TUMBLR_CONSUMER_SECRET', consumer_secret),
|
|
('$TUMBLR_OAUTH_TOKEN', oauth_token),
|
|
('$TUMBLR_OAUTH_SECRET', oauth_secret)] if val is None]
|
|
|
|
if missing_vars:
|
|
print("Missing important environment variables:", missing_vars)
|
|
sys.exit(1)
|
|
|
|
return pytumblr.TumblrRestClient(
|
|
consumer_key=consumer_key, # type: ignore
|
|
consumer_secret=consumer_secret, # type: ignore
|
|
oauth_token=oauth_token, # type: ignore
|
|
oauth_secret=oauth_secret, # type: ignore
|
|
)
|
|
|
|
|
|
def json_encoder(obj):
|
|
if dataclasses.is_dataclass(obj):
|
|
return dataclasses.asdict(obj) # type: ignore
|
|
elif isinstance(obj, pl.DataFrame):
|
|
return json.loads(obj.write_json())
|
|
raise TypeError(f"Type {type(obj)} is not serializable!")
|
|
|
|
|
|
def file_is_used(obj_arg: str, args: Dict[str, Any]) -> bool:
|
|
return obj_arg in args and args[obj_arg]
|
|
|
|
|
|
def load_from_file(path: str) -> pl.DataFrame:
|
|
return pl.read_json(source=Path(path), infer_schema_length=None)
|
|
|
|
|
|
def main() -> None:
|
|
args: Dict[str, Any] = get_args()
|
|
suppress = args['suppress']
|
|
client: pytumblr.TumblrRestClient = init_client()
|
|
|
|
# Read queued posts from a file, or fetch them via API and write queued post JSON.
|
|
queued_posts: pl.DataFrame = pl.DataFrame()
|
|
if not suppress or 'queued' not in suppress:
|
|
if file_is_used('queued_file', args):
|
|
queued_posts = load_from_file(args['queued_file'])
|
|
else:
|
|
queued_posts = fetcher.fetch_queued_posts(client, args)
|
|
queued_posts.write_json("queued_posts.json")
|
|
|
|
# Read drafts from a file, or fetch them via API and write draft JSON.
|
|
drafts: pl.DataFrame = pl.DataFrame()
|
|
if not suppress or 'drafts' not in suppress:
|
|
if file_is_used('drafts_file', args):
|
|
drafts = load_from_file(args['drafts_file'])
|
|
else:
|
|
drafts = fetcher.fetch_drafts(client, args)
|
|
drafts.write_json("drafts.json")
|
|
|
|
# Read posts from a file, or fetch them via API and write post JSON.
|
|
posts: pl.DataFrame = pl.DataFrame()
|
|
if not suppress or 'posts' not in suppress:
|
|
if file_is_used('posts_file', args):
|
|
posts = load_from_file(args['posts_file'])
|
|
else:
|
|
posts = fetcher.fetch_posts(client, args)
|
|
posts.write_json("posts.json")
|
|
|
|
# Read activity from a file, or fetch them via API and write activity JSON.
|
|
activity: pl.DataFrame = pl.DataFrame(infer_schema_length=None)
|
|
if not suppress or 'activity' not in suppress:
|
|
if file_is_used('activity_file', args):
|
|
activity = load_from_file(args['activity_file'])
|
|
else:
|
|
activity = fetcher.fetch_activity(client, args, posts)
|
|
activity.write_json('activity.json')
|
|
|
|
# The model.
|
|
model: StatsModel = StatsModel(blog_name=args['blog'],
|
|
posts = posts,
|
|
queued_posts = queued_posts,
|
|
drafts = drafts,
|
|
activity = activity,
|
|
tags = args['tags'])
|
|
|
|
# Write the model as JSON output.
|
|
# https://stackoverflow.com/q/26244323
|
|
with open('./tumblr_stats.json', 'w') as f:
|
|
json.dump(asdict(model), f, indent=1, default=json_encoder)
|
|
|
|
return
|
|
|
|
|
|
# DO NOT DELETE. The main if statement.
|
|
if __name__ == '__main__':
|
|
main()
|
|
print('All done.')
|
|
sys.exit(0)
|