diff --git a/src/raw_collect/twitter.py b/src/raw_collect/twitter.py index 9b9c518..4f5d62b 100644 --- a/src/raw_collect/twitter.py +++ b/src/raw_collect/twitter.py @@ -1,6 +1,5 @@ import json import math -import os import random import time from dataclasses import dataclass @@ -12,7 +11,7 @@ import pytz import tweepy from tweepy import API, TooManyRequests -from raw_collect.utils import Config, debug, Posting, json_stringify, load_config +from utils import Config, debug, Posting, json_stringify, load_config, normalize_directory @dataclass @@ -225,12 +224,7 @@ def download_users_execute(api: API, n: float, base_dir: str, rate_limit: int, :param next_set: The next set of starting users :return: None """ - - # Ensure that basedir doesn't ends with / - if base_dir == '': - base_dir = '.' - if base_dir.endswith('/'): - base_dir = base_dir[:-1] + base_dir = normalize_directory(base_dir) # Ensure directory exists Path(f'{base_dir}/users').mkdir(parents=True, exist_ok=True) diff --git a/src/raw_collect/twitter_individual.py b/src/raw_collect/twitter_individual.py index e73e9b0..a5f3617 100644 --- a/src/raw_collect/twitter_individual.py +++ b/src/raw_collect/twitter_individual.py @@ -1,5 +1,5 @@ from raw_collect.twitter import tweepy_login, download_user_tweets -from raw_collect.utils import * +from utils import * if __name__ == '__main__': diff --git a/src/raw_collect/utils.py b/src/utils.py similarity index 80% rename from src/raw_collect/utils.py rename to src/utils.py index 9b2b7c3..f29f409 100644 --- a/src/raw_collect/utils.py +++ b/src/utils.py @@ -56,6 +56,28 @@ def debug(msg: object) -> None: print('[DEBUG] ' + str(msg)) +def normalize_directory(directory: str) -> str: + """ + Normalize a directory input: Ensure that the directory doesn't end with "/", and ensure that an + empty directory input will be relative (".") + + >>> normalize_directory('') + '.' + >>> normalize_directory('path/') + 'path' + >>> normalize_directory('path') + 'path' + + :param directory: Input directory + :return: Normalized directory + """ + if directory == '': + directory = '.' + if directory.endswith('/'): + directory = directory[:-1] + return directory + + class EnhancedJSONEncoder(json.JSONEncoder): def default(self, o):