Skip to content

Commit

Permalink
add missing type hints for return values
Browse files Browse the repository at this point in the history
  • Loading branch information
redvox committed Jan 20, 2025
1 parent 0002f89 commit b75f6c5
Show file tree
Hide file tree
Showing 8 changed files with 40 additions and 43 deletions.
4 changes: 2 additions & 2 deletions app/aws/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ def fetch_role_credentials(user_name: str, profile_group: ProfileGroup) -> Resul
return result


def _remove_unused_profiles(credentials_file, profile_group: ProfileGroup):
def _remove_unused_profiles(credentials_file, profile_group: ProfileGroup) -> configparser:
used_profiles = profile_group.list_profile_names()
for profile in credentials_file.sections():
if profile not in used_profiles and \
Expand Down Expand Up @@ -228,7 +228,7 @@ def write_profile_config(profile_group: ProfileGroup, region: str) -> Result:
return result


def _remove_unused_configs(config_file: configparser, profile_group: ProfileGroup):
def _remove_unused_configs(config_file: configparser, profile_group: ProfileGroup) -> configparser:
used_profiles = profile_group.list_profile_names()

for config_name in config_file.sections():
Expand Down
32 changes: 16 additions & 16 deletions app/core/config.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from mimetypes import inited
from typing import List, Dict
from typing import List, Dict, Optional

from app.core import files
from app.core.profile_group import ProfileGroup
Expand All @@ -18,7 +18,7 @@ def __init__(self):
self.mfa_shell_command = None
self.default_access_key = None

def initialize(self):
def initialize(self) -> None:
config = files.load_config()
self.mfa_shell_command = config.get('mfa_shell_command', None)
self.default_access_key = config.get('default_access_key', _default_access_key)
Expand All @@ -29,7 +29,7 @@ def initialize(self):
self.initialize_profile_groups(accounts=accounts, service_roles=self.service_roles,
default_access_key=self.default_access_key)

def initialize_profile_groups(self, accounts: dict, service_roles: dict, default_access_key: str):
def initialize_profile_groups(self, accounts: dict, service_roles: dict, default_access_key: str) -> None:
self.profile_groups = {}
for group_name, group_data in accounts.items():
profile_group = ProfileGroup(name=group_name,
Expand All @@ -46,22 +46,22 @@ def initialize_profile_groups(self, accounts: dict, service_roles: dict, default
role_name=selected_service_role)
self.validate()

def set_mfa_shell_command(self, mfa_shell_command: str):
def set_mfa_shell_command(self, mfa_shell_command: str) -> None:
self.mfa_shell_command = mfa_shell_command

def set_default_access_key(self, default_access_key: str):
def set_default_access_key(self, default_access_key: str) -> None:
self.default_access_key = default_access_key

def save_config(self):
def save_config(self) -> None:
files.save_config_file({
'mfa_shell_command': self.mfa_shell_command,
'default_access_key': self.default_access_key,
})

def save_accounts(self):
def save_accounts(self) -> None:
files.save_accounts_file(self.to_dict())

def save_selected_service_role(self, group_name: str, profile_name: str, role_name: str):
def save_selected_service_role(self, group_name: str, profile_name: str, role_name: str) -> None:
if group_name not in self.service_roles:
self.service_roles[group_name] = {
'selected_profile': None,
Expand All @@ -75,7 +75,7 @@ def save_selected_service_role(self, group_name: str, profile_name: str, role_na
self.service_roles[group_name]['history'] = history
files.save_service_roles_file(self.service_roles)

def save_available_service_roles(self, group_name: str, profile_name: str, role_list: List[str]):
def save_available_service_roles(self, group_name: str, profile_name: str, role_list: List[str]) -> None:
if group_name not in self.service_roles:
self.service_roles[group_name] = {
'selected_profile': None,
Expand All @@ -86,20 +86,20 @@ def save_available_service_roles(self, group_name: str, profile_name: str, role_
self.service_roles[group_name]['available'][profile_name] = role_list
files.save_service_roles_file(self.service_roles)

def get_selected_service_role_source_profile(self, group: str):
def get_selected_service_role_source_profile(self, group: str) -> Optional[str]:
return self.service_roles.get(group, {}).get('selected_profile')

def get_selected_service_role(self, group: str):
def get_selected_service_role(self, group: str) -> Optional[str]:
return self.service_roles.get(group, {}).get('selected_role')

def get_available_service_roles(self, group: str, profile: str):
def get_available_service_roles(self, group: str, profile: str) -> List[str]:
return self.service_roles.get(group, {}).get('available', {}).get(profile, [])

def get_history(self, group: str):
return self.service_roles.get(group, {}).get('history', [])

@staticmethod
def _add_to_history(profile: str, role: str, history: List[str]):
def _add_to_history(profile: str, role: str, history: List[str]) -> List[str]:
if not profile or not role:
return history
history_string = f'{profile} : {role}'
Expand All @@ -125,13 +125,13 @@ def validate(self) -> None:
self.valid = valid
self.error = error

def list_groups(self):
def list_groups(self) -> List[ProfileGroup]:
return list(self.profile_groups.values())

def get_group(self, name):
def get_group(self, name) -> Optional[ProfileGroup]:
return self.profile_groups.get(name, None)

def to_dict(self):
def to_dict(self) -> dict:
d = {}
for name, group in self.profile_groups.items():
d[name] = group.to_dict()
Expand Down
12 changes: 6 additions & 6 deletions app/core/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def login_gcp(self, profile_group: ProfileGroup) -> Result:
result.set_success()
return result

def logout(self):
def logout(self) -> Result:
result = Result()
logger.info(f'start logout')
self.active_profile_group = None
Expand Down Expand Up @@ -129,10 +129,10 @@ def get_region(self) -> Optional[str]:
return self.active_profile_group.region
return None

def get_profile_group_list(self):
def get_profile_group_list(self) -> List[ProfileGroup]:
return self.config.list_groups()

def get_active_profile_color(self):
def get_active_profile_color(self) -> str:
return self.active_profile_group.color

def rotate_access_key(self, access_key: str, mfa_token: Optional[str]) -> Result:
Expand Down Expand Up @@ -197,7 +197,7 @@ def set_service_role(self, profile_name: str, role_name: str) -> Result:
result.set_success()
return result

def set_available_service_roles(self, profile, role_list: List[str]):
def set_available_service_roles(self, profile, role_list: List[str]) -> Result:
result = Result()
logger.info('set available service roles')
self.config.save_available_service_roles(group_name=self.active_profile_group.name,
Expand All @@ -207,7 +207,7 @@ def set_available_service_roles(self, profile, role_list: List[str]):
return result

@staticmethod
def run_script(profile_group: ProfileGroup):
def run_script(profile_group: ProfileGroup) -> Result:
result = Result()
if not profile_group or not profile_group.script:
result.set_success()
Expand Down Expand Up @@ -252,5 +252,5 @@ def set_access_key(key_name: str, key_id: str, key_secret: str) -> Result:
return credentials.set_access_key(key_name=key_name, key_id=key_id, key_secret=key_secret)

@staticmethod
def get_access_key_list():
def get_access_key_list() -> list:
return credentials.get_access_key_list()
28 changes: 14 additions & 14 deletions app/core/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,70 +47,70 @@ def get_active_group_file_path() -> str:
return f'{get_app_path()}/{active_group_file_name}'


def parse_yaml(text: str):
def parse_yaml(text: str) -> dict:
try:
return yamli.load(text) or {}
except (ParserError, ScannerError):
return {}


def dump_yaml(d: dict):
def dump_yaml(d: dict) -> str:
buffer = io.BytesIO()
yamli.dump(d, buffer)
return buffer.getvalue().decode("utf-8")


def _load_file(path):
def _load_file(path) -> str:
if os.path.exists(path):
with open(path, 'r') as file:
return file.read()
return ''


def _write_file(path, content):
def _write_file(path, content) -> None:
with open(path, 'w') as file:
file.write(str(content))


def remove_file(path):
def remove_file(path) -> None:
if os.path.exists(path):
os.remove(path)


def file_exists(file_path):
def file_exists(file_path) -> bool:
file_path = file_path.strip()
if ' ' in file_path:
file_path = file_path.split(' ')[0]
return os.path.exists(file_path)


def load_config():
def load_config() -> dict:
return parse_yaml(_load_file(get_config_path()))


def load_accounts():
def load_accounts() -> dict:
return parse_yaml(_load_file(get_accounts_path()))


def load_service_roles():
def load_service_roles() -> dict:
return parse_yaml(_load_file(get_service_roles_path()))


def save_config_file(config_dict: dict):
def save_config_file(config_dict: dict) -> None:
_write_file(get_config_path(), dump_yaml(config_dict))


def save_accounts_file(account_dict: dict):
def save_accounts_file(account_dict: dict) -> None:
_write_file(get_accounts_path(), dump_yaml(account_dict))


def save_service_roles_file(service_roles: dict):
def save_service_roles_file(service_roles: dict) -> None:
_write_file(get_service_roles_path(), dump_yaml(service_roles))


def load_logs():
def load_logs() -> str:
return _load_file(get_log_path()) or 'no logs found'


def write_active_group_file(group_name):
def write_active_group_file(group_name) -> None:
_write_file(get_active_group_file_path(), group_name)
2 changes: 1 addition & 1 deletion app/core/profile.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def validate(self) -> (bool, str):
return False, f'a profile in {self.group.name} has no role'
return True, ''

def to_dict(self):
def to_dict(self) -> dict:
d = {
'profile': self.profile,
'account': self.account,
Expand Down
2 changes: 1 addition & 1 deletion app/core/profile_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def get_access_key(self) -> str:
return self.access_key
return self.default_access_key

def set_service_role_profile(self, source_profile_name, role_name):
def set_service_role_profile(self, source_profile_name, role_name) -> None:
source_profile = self.get_profile(profile_name=source_profile_name)
if not source_profile:
logger.warning(f'source profile {source_profile_name} not found. Unset service_profile.')
Expand Down
2 changes: 0 additions & 2 deletions app/gcp/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,5 @@ def set_default_quota_project(project: str):


def set_default_region(region: str):

logger.info(f'set default region to: {region}')
return shell.run(f"gcloud config set compute/region {region}", timeout=5)

1 change: 0 additions & 1 deletion app/gui/trayicon.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from functools import partial
from typing import List, TYPE_CHECKING

from PyQt6.QtCore import QSize
from PyQt6.QtWidgets import QSystemTrayIcon, QMenu

from app.aws import regions
Expand Down

0 comments on commit b75f6c5

Please sign in to comment.