ensemble_machine / save_data.py
Naisong Zhou
data save position fixed to id-determined row
d089d7b
import os
from google.oauth2 import service_account
from googleapiclient.discovery import build
import json
import os
def load_envs(local=False):
"""Load the environment variables."""
if local:
from dotenv import load_dotenv
load_dotenv(override=True)
service_account_info = json.loads(os.environ['GOOGLE_APPLICATION_CREDENTIALS_JSON'])
SHEET_ID1 = os.environ['SHEET_ID1'] # human-ai-sequential
SHEET_ID2 = os.environ['SHEET_ID2'] # human-ai-parallel
SHEET_ID3 = os.environ['SHEET_ID3'] # ai-human-sequential
SHEET_IDs = [SHEET_ID1, SHEET_ID2, SHEET_ID3]
return service_account_info, SHEET_IDs
def get_sheet_service(local=False):
"""Get the google sheet service object."""
service_account_info, SHEET_IDs = load_envs(local=local)
# verify the service_account_info
credentials = service_account.Credentials.from_service_account_info(
service_account_info,
scopes=['https://www.googleapis.com/auth/spreadsheets']
)
# build the service object
service = build('sheets', 'v4', credentials=credentials)
return service, SHEET_IDs
def col_letter(col_num):
"""Convert a column number to its corresponding Excel-style letter."""
letter = ''
while col_num > 0:
col_num, remainder = divmod(col_num - 1, 26)
letter = chr(65 + remainder) + letter
return letter
def add_or_update_row_at_fixed_position(row_id, new_row, service, SPREADSHEET_ID, num_of_columns=5):
"""
Add or update data at a fixed row position in Google Sheet.
Args:
row_id: int, the row number (1-based index) where the data should be written.
new_row: list, the new data to be added or updated.
service: Google Sheets API service instance.
SPREADSHEET_ID: str, the spreadsheet ID.
num_of_columns: int, the number of columns to consider (default: 5).
"""
# Ensure `new_row` has exactly `num_of_columns` elements
new_row = new_row + [""] * (num_of_columns - len(new_row)) if len(new_row) < num_of_columns else new_row[:num_of_columns]
# Determine the write range for the specified row
range_to_write = f"Sheet1!A{row_id}:{col_letter(num_of_columns)}{row_id}"
# Directly write the data to the specified row
response = service.spreadsheets().values().update(
spreadsheetId=SPREADSHEET_ID,
range=range_to_write,
valueInputOption='RAW',
body={"values": [new_row]}
).execute()
print(f"Data written to row {row_id}. Response: {response}")
if __name__ == "__main__":
service, SHEET_IDs = get_sheet_service(local=True)
new_row = ["test1", "test2", "test3", "test4", "test5"]
add_or_update_row_at_fixed_position(12, new_row, service, SHEET_IDs[0])