oauth handler

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
+__all__ = ['GoogleDocsConnector', 'DocumentGenerator', 'TemplateCreator', 'GoogleAPIError', 
+           'DocumentGenerationError', 'TemplateCreationError', 'TemplateParser', 
+           'InvalidTemplateException', 'OAuthHandler']
+__version__ = '0.2.0'

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.

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