diff --git a/src/doctemplify/__init__.py b/src/doctemplify/__init__.py index d487c37..735a298 100644 --- a/src/doctemplify/__init__.py +++ b/src/doctemplify/__init__.py @@ -3,6 +3,9 @@ from .template_creator import TemplateCreator from .exceptions import GoogleAPIError, DocumentGenerationError, TemplateCreationError from .template_parser import TemplateParser, InvalidTemplateException +from .oauth_handler import OAuthHandler -__all__ = ['GoogleDocsConnector', 'DocumentGenerator', 'TemplateCreator', 'GoogleAPIError', 'DocumentGenerationError', 'TemplateCreationError', 'TemplateParser', 'InvalidTemplateException'] -__version__ = '0.1.1' # Increment the version number \ No newline at end of file +__all__ = ['GoogleDocsConnector', 'DocumentGenerator', 'TemplateCreator', 'GoogleAPIError', + 'DocumentGenerationError', 'TemplateCreationError', 'TemplateParser', + 'InvalidTemplateException', 'OAuthHandler'] +__version__ = '0.2.0' \ No newline at end of file diff --git a/src/doctemplify/google_docs_connector.py b/src/doctemplify/google_docs_connector.py index ef93e22..110617c 100644 --- a/src/doctemplify/google_docs_connector.py +++ b/src/doctemplify/google_docs_connector.py @@ -1,25 +1,42 @@ -from google.oauth2.service_account import Credentials +from google.oauth2.service_account import Credentials as ServiceAccountCredentials from googleapiclient.discovery import build from typing import Dict, Any, List, Tuple, Optional from .google_fonts import GOOGLE_DOCS_FONTS from .exceptions import GoogleAPIError from .template_parser import TemplateParser, InvalidTemplateException +from .oauth_handler import OAuthHandler GDOCS_DEFAULT_URL = "https://docs.google.com/document/d/{}/edit" class GoogleDocsConnector: - def __init__(self, service_account_file: str): + def __init__(self, auth_file: str, auth_type: str = 'service_account'): self.SCOPES = ['https://www.googleapis.com/auth/drive', 'https://www.googleapis.com/auth/documents'] self.GOOGLE_DOCS_FONTS = GOOGLE_DOCS_FONTS + self.auth_type = auth_type + self.auth_file = auth_file + self.drive_service = None + self.docs_service = None + self.template_parser = None + self.current_index = 1 + self._initialize_services() + + def _initialize_services(self): try: - self.creds = Credentials.from_service_account_file(service_account_file, scopes=self.SCOPES) - self.drive_service = build('drive', 'v3', credentials=self.creds) - self.docs_service = build('docs', 'v1', credentials=self.creds) + if self.auth_type == 'service_account': + creds = ServiceAccountCredentials.from_service_account_file(self.auth_file, scopes=self.SCOPES) + elif self.auth_type == 'oauth': + oauth_handler = OAuthHandler(self.auth_file, self.SCOPES) + creds = oauth_handler.get_credentials() + else: + raise ValueError(f"Invalid auth_type: {self.auth_type}. Must be 'service_account' or 'oauth'.") + + self.drive_service = build('drive', 'v3', credentials=creds) + self.docs_service = build('docs', 'v1', credentials=creds) self.template_parser = TemplateParser(self) - self.current_index = 1 except Exception as e: raise GoogleAPIError(f"Failed to initialize Google API services: {str(e)}") + def create_document(self, title: str, public: bool = False) -> Tuple[str, str]: """ Create a new Google Docs document. diff --git a/src/doctemplify/oauth_handler.py b/src/doctemplify/oauth_handler.py new file mode 100644 index 0000000..24c14c2 --- /dev/null +++ b/src/doctemplify/oauth_handler.py @@ -0,0 +1,29 @@ +import os +from google_auth_oauthlib.flow import InstalledAppFlow +from google.oauth2.credentials import Credentials +from google.auth.transport.requests import Request + +class OAuthHandler: + def __init__(self, client_secrets_file, scopes): + self.client_secrets_file = client_secrets_file + self.scopes = scopes + self.credentials = None + + def authenticate(self): + if os.path.exists('token.json'): + self.credentials = Credentials.from_authorized_user_file('token.json', self.scopes) + + if not self.credentials or not self.credentials.valid: + if self.credentials and self.credentials.expired and self.credentials.refresh_token: + self.credentials.refresh(Request()) + else: + flow = InstalledAppFlow.from_client_secrets_file(self.client_secrets_file, self.scopes) + self.credentials = flow.run_local_server(port=0) + + with open('token.json', 'w') as token: + token.write(self.credentials.to_json()) + + def get_credentials(self): + if not self.credentials: + self.authenticate() + return self.credentials \ No newline at end of file