File size: 3,074 Bytes
7a0c684 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
"""
Script to automatically update all files with token from .env.
"""
import os
import re
from pathlib import Path
from typing import List, Set
from dotenv import load_dotenv
def load_token() -> str:
"""Load token from .env file."""
root_dir = Path(__file__).parent.parent
env_path = root_dir / '.env'
load_dotenv(dotenv_path=env_path)
token = os.getenv('HF_TOKEN')
if not token:
raise ValueError("HF_TOKEN not set in .env file")
return token
def find_python_files(start_path: str) -> List[Path]:
"""Find all Python files in the directory tree."""
return [
Path(os.path.join(root, file))
for root, _, files in os.walk(start_path)
for file in files
if file.endswith('.py') and 'venv' not in root and '__pycache__' not in root
]
def has_get_token_cached(content: str) -> bool:
"""Check if file uses get_hf_token_cached."""
return (
'get_hf_token_cached' in content and
'from config import' in content
)
def update_file(file_path: Path, token: str) -> bool:
"""Update a single file with token."""
print(f"Checking {file_path}")
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
# Skip files that don't use the config token
if not has_get_token_cached(content):
print(f"Skipping {file_path} - no token config usage found")
return False
print(f"Updating {file_path} with token")
# Add token initialization right after config import
lines = content.split('\n')
modified_lines = []
for line in lines:
modified_lines.append(line)
if line.strip() == 'from config import get_hf_token_cached':
# Add token setup after import
modified_lines.append(f'\n# Initialize token from .env')
modified_lines.append(f'os.environ["HF_TOKEN"] = "{token}"\n')
content = '\n'.join(modified_lines)
# Write changes
with open(file_path, 'w', encoding='utf-8') as f:
f.write(content)
print(f"Updated {file_path}")
return True
def main():
"""Main function to update all files."""
try:
# Get token from .env
token = load_token()
# Get project root
root_dir = Path(__file__).parent.parent
# Find and update all Python files
updated = 0
python_files = find_python_files(root_dir)
for file_path in python_files:
try:
if update_file(file_path, token):
updated += 1
except Exception as e:
print(f"Error updating {file_path}: {e}")
print(f"\nUpdated {updated} files with token from .env")
except Exception as e:
print(f"Error: {e}")
print("Please make sure HF_TOKEN is set in your .env file")
return 1
return 0
if __name__ == '__main__':
main()
|