build_draft_stats, total_original_post_notes_by_qtr_and_year, other minor improvements

tumblr_stats.py

@@ -11,6 +11,7 @@ from typing import Any, Callable, Dict, List, Tuple
 
 import pytumblr
 
+from build_draft_stats_model import BuildDraftStatsModel
 from build_tag_stats_model import BuildTagStatsModel
 from build_total_stats_model import BuildTotalStatsModel
 from build_queue_stats_model import BuildQueueStatsModel
@@ -25,7 +26,7 @@ def get_args() -> Dict[str, Any]:
         + '$TUMBLR_CONSUMER_KEY, $TUMBLR_CONSUMER_SECRET, $TUMBLR_OAUTH_TOKEN, and $TUMBLR_OAUTH_SECRET',
         epilog='— Be gay and do crime')
     parser.add_argument('operation', type=str, nargs='+', metavar='OPERATION',
-                        choices=['build_tag_stats', 'build_queue_stats'],
+                        choices=['build_tag_stats', 'build_queue_stats', 'build_draft_stats'],
                         help="operation used to calculate stats")
     parser.add_argument('-b', '--blog', type=str, required=True,
                         help='blog name for which to calculate stats')
@@ -76,7 +77,6 @@ def filter_posts_for_after(post_list: List[Dict[str, Any]],
         x['date'], '%Y-%m-%d %H:%M:%S %Z') > after
     return [post for post in post_list if after_check(post)]
 
-
 def build_post_maps(client: pytumblr.TumblrRestClient,
                     args: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     og_post_map: Dict[str, Any] = {}
@@ -92,6 +92,9 @@ def build_post_maps(client: pytumblr.TumblrRestClient,
     #     before: datetime = args['before']
     #     params.update({'before': int(before.timestamp())})
 
+    draft_url = f"/v2/blog/{blog_name}/posts/draft"
+    is_draft_stats: bool = 'build_draft_stats' in args['operation']
+
     total: int = 0
     offset: int = 0
     limit: int = 20
@@ -101,12 +104,14 @@ def build_post_maps(client: pytumblr.TumblrRestClient,
         # Begin LOOP
         # Get me some posts via REST! 😈🍪🍪🍪
         data: Dict[str, Any]
-        if 'build_queue_stats' in args['operation'] and len(args['operation']) == 1:
+        if 'build_queue_stats' in args['operation']:
             data = client.queue(f"{blog_name}.tumblr.com",
                                 offset=offset,
                                 limit=limit,
                                 **params)
-        else:  # Above is for queued posts, below is for published posts.
+        elif is_draft_stats:
+            data = client.send_api_request("get", draft_url)
+        else:  # Above is for queued + draft posts, below is for published posts.
             data = client.posts(f"{blog_name}.tumblr.com",
                                 offset=offset,
                                 limit=limit,
@@ -119,7 +124,7 @@ def build_post_maps(client: pytumblr.TumblrRestClient,
             break
 
         next_off: int = 0
-        if '_links' in data:
+        if '_links' in data and not is_draft_stats:
             links = data['_links']
             if 'next' in links and 'query_params' in links['next']:
                 next_off = int(links['next']['query_params']['offset'])
@@ -135,7 +140,7 @@ def build_post_maps(client: pytumblr.TumblrRestClient,
             after: datetime = args['after']
             curr_posts = filter_posts_for_after(curr_posts, after)
             if not curr_posts:
-                print(f"All posts after {after.year}-{after.month} processed.")
+                print(f"All posts after {after.year}-{after.month}-{after.day} processed.")
                 return (og_post_map, un_og_post_map)
 
         # This block populates the local post_maps from the raw response data.
@@ -150,6 +155,15 @@ def build_post_maps(client: pytumblr.TumblrRestClient,
         og_post_map.update(local_og_post_map)
         un_og_post_map.update(local_un_og_post_map)
 
+        # For build_draft_stats.
+        if is_draft_stats:
+            if '_links' in data:
+                draft_url = data['_links']['next']['href']
+                continue
+            else:
+                print('All draft posts processed.')
+                break
+
         # The increment and status printing.
         if next_off != 0 and next_off != offset:
             offset = next_off
@@ -167,6 +181,13 @@ def build_post_maps(client: pytumblr.TumblrRestClient,
 def main() -> None:
     args: Dict[str, Any] = get_args()
     client: pytumblr.TumblrRestClient = init_client()
+    operation: List[str] = args['operation']
+
+    # Quick bail for bad use of build_queue_stats and build_draft_stats.
+    if (set(operation) & set(['build_queue_stats', 'build_draft_stats'])) and len(operation) > 1:
+        print(f"You can't mix operations {operation} together. Sorry.")
+        sys.exit(1)
+        pass
 
     # Handle JSON input (if you don't want to make API calls.)
     if 'input' in args and args['input']:
@@ -193,28 +214,33 @@ def main() -> None:
 
     # Pick a stats model, which will determine output.
     stats_model: StatsModel
-    if 'build_queue_stats' in args['operation']:
-        if len(args['operation']) != 1:
-            print('You can\'t mix build_queue_stats with other operations. Sorry.')
+    match args:
+        case {'operation': ['build_queue_stats']}:
+            stats_model = BuildQueueStatsModel(blog_name=args['blog'],
+                                            original_post_map=og_post_map,
+                                            unoriginal_post_map=un_og_post_map)
+        case {'operation': ['build_draft_stats']}:
+            stats_model = BuildDraftStatsModel(blog_name=args['blog'],
+                                            original_post_map=og_post_map,
+                                            unoriginal_post_map=un_og_post_map)
+        case {'operation': op} if 'build_tag_stats' in operation:
+            stats_model = BuildTagStatsModel(blog_name=args['blog'],
+                                            original_post_map=og_post_map,
+                                            unoriginal_post_map=un_og_post_map)
+            stats_model.tags = args['tags']
+        case {'operation': op} if 'build_total_stats' in operation:
+            if 'before' not in args:  # or 'after' not in args:
+                print(f"You must specify a time range for {op}. " +
+                    'You\'ll otherwise request TOO MUCH DATA!')
+                sys.exit(1)
+            stats_model = BuildTotalStatsModel(blog_name=args['blog'],
+                                            original_post_map=og_post_map,
+                                            unoriginal_post_map=un_og_post_map)
+        case _:
+            print('Unsupported command. How did you even make it this far?!')
             sys.exit(1)
-        stats_model = BuildQueueStatsModel(blog_name=args['blog'],
-                                           original_post_map=og_post_map,
-                                           unoriginal_post_map=un_og_post_map)
-    if 'build_tag_stats' in args['operation']:
-        stats_model = BuildTagStatsModel(blog_name=args['blog'],
-                                         original_post_map=og_post_map,
-                                         unoriginal_post_map=un_og_post_map)
-        stats_model.tags = args['tags']
-    if 'build_total_stats' in args['operation']:
-        if 'before' not in args:  # or 'after' not in args:
-            print('You must specify a time range for build_total stats. ' +
-                  'You\'ll otherwise request TOO MUCH DATA!')
-            sys.exit(1)
-        stats_model = BuildTotalStatsModel(blog_name=args['blog'],
-                                           original_post_map=og_post_map,
-                                           unoriginal_post_map=un_og_post_map)
 
-    # Write the chosen model as JSON output.
+    # Write the selected model as JSON output.
     with open('./tumblr_stats.json', 'w') as f:
         json.dump(asdict(stats_model), f, indent=1, default=str)