Coverage for src/prepare_times_nz/utilities/helpers.py: 32%

1""" 

2This script currently contains miscellaneous functions 

3 

4THese may need to be organised more clearly within the program, as there 

5is not a clear conceptual throughline. 

6 

7It currently includes: 

8 

9Directory manipulation functions: 

10 - clear_data_intermediate() 

11 - clear_output() 

12 

13Data testing functions: 

14 - check_table_grain() 

15 

16Data manipulation functions: 

17 - select_and_rename() 

18 

19These should potentially be moved into their own specific scripts, 

20OR this might make more sense if we can move more functions to shared 

21scripts with clearer conceptual throughlines. TBD. 

22 

23""" 

24 

25import os 

26import shutil 

27 

28from prepare_times_nz.utilities.filepaths import DATA_INTERMEDIATE, OUTPUT_LOCATION 

29from prepare_times_nz.utilities.logger_setup import logger 

30 

31 

32def clear_data_intermediate(): 

33 """ 

34 DELETES the folder defined as DATA_INTERMEDIATE 

35 """ 

36 # Delete folder 

37 if os.path.exists(DATA_INTERMEDIATE): 

38 logger.debug("DATA_INTERMEDIATE = {%s}", DATA_INTERMEDIATE) 

39 shutil.rmtree(DATA_INTERMEDIATE) 

40 # and make fresh 

41 os.makedirs(DATA_INTERMEDIATE) 

42 

43 

44def clear_output(): 

45 """ 

46 DELETES the folder defined as OUTPUT_LOCATION 

47 """ 

48 # Delete folder 

49 if os.path.exists(OUTPUT_LOCATION): 

50 logger.debug("OUTPUT_LOCATION = {%s}", OUTPUT_LOCATION) 

51 shutil.rmtree(OUTPUT_LOCATION) 

52 # and make fresh 

53 os.makedirs(OUTPUT_LOCATION) 

54 

55 

56# Some data manipulation functions 

57 

58 

59def select_and_rename(df, name_map): 

60 """ 

61 Selects and renames columns in a DataFrame based on a provided mapping. 

62 

63 Parameters: 

64 - df: The input DataFrame. 

65 - name_map: A dictionary: 

66 keys are the original column names and values are the new column names. 

67 

68 Returns: 

69 - A DataFrame with selected and renamed columns. 

70 """ 

71 # Select columns based on the mapping 

72 selected_df = df[list(name_map.keys())].copy() 

73 

74 # Rename columns 

75 selected_df.rename(columns=name_map, inplace=True) 

76 

77 return selected_df 

78 

79 

80# Some tests 

81 

82 

83def check_table_grain(df, grain_list): 

84 """ 

85 Check the grain of a DataFrame against a list of expected grains. 

86 

87 Parameters: 

88 - df: The input DataFrame. 

89 - grain_list: A list of column names 

90 These should uniquely identify the rows in the Dataframe. 

91 

92 Returns: 

93 - A boolean indicating whether rows are uniquely idenfied by the grain_list 

94 """ 

95 return df.duplicated(subset=grain_list).sum() == 0 

96 

97 

98def test_table_grain(df, grain_list): 

99 """ 

100 A wrapper and logging output for check_table_grain 

101 

102 Parameters: 

103 - df: The input DataFrame. 

104 - grain_list: A list of column names that should uniquely identify the rows in df. 

105 

106 Outputs: logging information regarding the results of the test. 

107 

108 """ 

109 

110 if check_table_grain(df, grain_list): 

111 logger.info( 

112 "Success: rows are uniquely identified using the following variables:" 

113 ) 

114 for var in grain_list: 

115 logger.info(" - {%s}", var) 

116 

117 else: 

118 logger.warning( 

119 "Please review! Rows are NOT uniquely identified using these variables:" 

120 ) 

121 for var in grain_list: 

122 logger.info(" - {%s}", var)