aail-hf/ensemble_machine
0
1import os2from google.oauth2 import service_account3from googleapiclient.discovery import build4import json5 6 7import os8 9def load_envs(local=False):10 """Load the environment variables."""11 if local:12 from dotenv import load_dotenv13 load_dotenv(override=True)14 service_account_info = json.loads(os.environ['GOOGLE_APPLICATION_CREDENTIALS_JSON'])15 SHEET_ID1 = os.environ['SHEET_ID1'] # human-ai-sequential16 SHEET_ID2 = os.environ['SHEET_ID2'] # human-ai-parallel17 SHEET_ID3 = os.environ['SHEET_ID3'] # ai-human-sequential18 SHEET_IDs = [SHEET_ID1, SHEET_ID2, SHEET_ID3]19 return service_account_info, SHEET_IDs20 21 22def get_sheet_service(local=False):23 """Get the google sheet service object."""24 service_account_info, SHEET_IDs = load_envs(local=local)25 # verify the service_account_info26 credentials = service_account.Credentials.from_service_account_info(27 service_account_info,28 scopes=['https://www.googleapis.com/auth/spreadsheets']29 )30 31 # build the service object32 service = build('sheets', 'v4', credentials=credentials)33 return service, SHEET_IDs34 35 36def col_letter(col_num):37 """Convert a column number to its corresponding Excel-style letter."""38 letter = ''39 while col_num > 0:40 col_num, remainder = divmod(col_num - 1, 26)41 letter = chr(65 + remainder) + letter42 return letter43 44def add_or_update_row_at_fixed_position(row_id, new_row, service, SPREADSHEET_ID, num_of_columns=5):45 """46 Add or update data at a fixed row position in Google Sheet.47 48 Args:49 row_id: int, the row number (1-based index) where the data should be written.50 new_row: list, the new data to be added or updated.51 service: Google Sheets API service instance.52 SPREADSHEET_ID: str, the spreadsheet ID.53 num_of_columns: int, the number of columns to consider (default: 5).54 """55 # Ensure `new_row` has exactly `num_of_columns` elements56 new_row = new_row + [""] * (num_of_columns - len(new_row)) if len(new_row) < num_of_columns else new_row[:num_of_columns]57 58 # Determine the write range for the specified row59 range_to_write = f"Sheet1!A{row_id}:{col_letter(num_of_columns)}{row_id}"60 61 # Directly write the data to the specified row62 response = service.spreadsheets().values().update(63 spreadsheetId=SPREADSHEET_ID,64 range=range_to_write,65 valueInputOption='RAW',66 body={"values": [new_row]}67 ).execute()68 69 print(f"Data written to row {row_id}. Response: {response}")70 71 72if __name__ == "__main__":73 service, SHEET_IDs = get_sheet_service(local=True)74 new_row = ["test1", "test2", "test3", "test4", "test5"]75 add_or_update_row_at_fixed_position(12, new_row, service, SHEET_IDs[0])76 77 