diff --git a/abr-testing/abr_testing/automation/google_sheets_tool.py b/abr-testing/abr_testing/automation/google_sheets_tool.py index 0a2e2bdfeb5..e4b8b6499e0 100644 --- a/abr-testing/abr_testing/automation/google_sheets_tool.py +++ b/abr-testing/abr_testing/automation/google_sheets_tool.py @@ -48,12 +48,13 @@ def open_worksheet(self, tab_number: int) -> Any: """Open individual worksheet within a googlesheet.""" return self.spread_sheet.get_worksheet(tab_number) - def create_worksheet(self, tab_name: int) -> None: + def create_worksheet(self, title: str) -> None: """Create a worksheet with tab name. Existing spreadsheet needed.""" try: - self.spread_sheet.add_worksheet(tab_name, rows="1000", cols="26") + new_sheet = self.spread_sheet.add_worksheet(title, rows="2000", cols="26") + return new_sheet.id except gspread.exceptions.APIError: - print("Work Sheet already exists") + print("Sheet already exists.") def write_header(self, header: List) -> None: """Write Header to first row if not present.""" @@ -61,7 +62,7 @@ def write_header(self, header: List) -> None: if header_list != header: self.worksheet.insert_row(header, self.row_index) - def write_to_row(self, data: List) -> None: + def write_to_row(self, data: List, title: str = "Sheet1") -> None: """Write data into a row in a List[] format.""" try: self.row_index += 1 @@ -106,11 +107,24 @@ def batch_delete_rows(self, row_indices: List[int]) -> None: } self.spread_sheet.batch_update(body=delete_body) + def batch_update_cells( + self, sheet_title: str, data: List[List[str]], start_column: str, start_row: int + ) -> None: + """Writes to multiple cells at once in a specific sheet.""" + sheet = self.spread_sheet.worksheet(sheet_title) + for idx, values in enumerate(data): + column = chr(ord(start_column) + idx) # Convert index to column letter + location = f"{column}{start_row}:{column}{start_row + len(values) - 1}" + cells_to_update = sheet.range(location) + for cell, value in zip(cells_to_update, values): + cell.value = value + sheet.update_cells(cells_to_update) + def update_cell( - self, row: int, column: int, single_data: Any + self, sheet_title: str, row: int, column: int, single_data: Any ) -> Tuple[int, int, Any]: """Update ONE individual cell according to a row and column.""" - self.worksheet.update_cell(row, column, single_data) + self.spread_sheet.worksheet(sheet_title).update_cell(row, column, single_data) return row, column, single_data def get_all_data(self) -> List[Dict[str, Any]]: @@ -121,6 +135,15 @@ def get_column(self, column_number: int) -> Set[str]: """Get all values in column.""" return self.worksheet.col_values(column_number) + def get_cell(self, cell: str) -> Any: + """Get cell value with location ex: 'A1'.""" + return self.worksheet.acell(cell).value + + def get_single_col_range(self, range: str) -> List: + """Get cell values from one column range.""" + values_range = self.worksheet.range(range) + return [cell.value for cell in values_range] + def get_index_row(self) -> int: """Check for the next available row to write too.""" row_index = len(self.get_column(1)) @@ -166,6 +189,9 @@ def create_line_chart( titles: List[str], series: List[Dict[str, Any]], domains: List[Dict[str, Any]], + axis: Dict[str, Any], + col_position: int = 0, + sheet_id: str = "0", ) -> None: """Create chart of data on google sheet.""" request_body = { @@ -178,10 +204,7 @@ def create_line_chart( "basicChart": { "chartType": "LINE", "legendPosition": "RIGHT_LEGEND", - "axis": [ - {"position": "BOTTOM_AXIS", "title": titles[1]}, - {"position": "LEFT_AXIS", "title": titles[2]}, - ], + "axis": axis, "domains": domains, "series": series, "headerCount": 1, @@ -190,9 +213,9 @@ def create_line_chart( "position": { "overlayPosition": { "anchorCell": { - "sheetId": 0, + "sheetId": sheet_id, "rowIndex": 1, - "columnIndex": 1, + "columnIndex": col_position, } } }, diff --git a/abr-testing/abr_testing/data_collection/abr_lpc.py b/abr-testing/abr_testing/tools/abr_lpc.py similarity index 100% rename from abr-testing/abr_testing/data_collection/abr_lpc.py rename to abr-testing/abr_testing/tools/abr_lpc.py