Coverage for src\gtrends_collection\utils.py: 92%

1"""Utility function to interact with framework resources and trends data.""" 

2 

3from os.path import isfile 

4from typing import Dict, List 

5 

6import pandas 

7 

8 

9def read_scope(scope_dir: str, which: str) -> List[str]: 

10 """ 

11 Reads in a scope file. 

12 

13 Args: 

14 scope_dir (str): Directory containing scope files. 

15 which (str): Name of the file to be read in (e.g., `locations`). 

16 

17 Examples: 

18 ```python 

19 terms = read_scope("./scope", "terms") 

20 ``` 

21 

22 Returns: 

23 A list of terms or locations. 

24 """ 

25 with open(f"{scope_dir}/{which}.txt", encoding="utf-8") as content: 

26 locations = [code.strip() for code in content.readlines()] 

27 return locations 

28 

29 

30def full_term_names(scope_dir: str, terms: List[str], include_id: bool = True) -> List[str]: 

31 """ 

32 Converts topic and category IDs to their full names, based on `scope_dir/term_map.csv`. 

33 

34 Args: 

35 terms (List[str]): Terms to be converted. 

36 

37 Examples: 

38 ```python 

39 full_term_names("./scope", ["/m/0cycc", "/g/11hy9m64ws"]) 

40 ``` 

41 

42 Returns: 

43 A version of `terns` with any matching terms converted. 

44 """ 

45 map_file = f"{scope_dir}/term_map.csv" 

46 if not isfile(map_file): 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true

47 return terms 

48 term_map = pandas.read_csv(map_file, index_col="id")["name"].to_dict() 

49 return [term_map.get(term, term) + (f" ({term})" if include_id and term.startswith("/") else "") for term in terms] 

50 

51 

52def full_metro_area_codes(scope_dir: str, locations: List[str]) -> List[str]: 

53 """ 

54 Adds country and state codes to metro area codes (e.g., `630` becomes `US-AL-630`), 

55 based on `scope_dir/locations.txt`. 

56 

57 Args: 

58 locations (List[str]): Locations to potentially prepend full location codes to. 

59 

60 Examples: 

61 ```python 

62 full_metro_area_codes("./scope", ["630", "743"]) 

63 ``` 

64 

65 Returns: 

66 A version of `locations` with any matching locations expanded. 

67 """ 

68 location_map: Dict[str, str] = {} 

69 for location in read_scope(scope_dir, "locations"): 

70 if len(location) == 9: 

71 location_parts = location.split("-") 

72 location_map[location_parts[2]] = location 

73 return [location_map.get(loc, loc) for loc in locations]