1# Copyright 2024, The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15"""Classes used to handle banners.""" 16 17from __future__ import annotations 18 19from datetime import date 20import json 21import logging 22from pathlib import Path 23from typing import Any, Callable 24 25from atest import atest_utils 26from atest import constants 27 28 29class BannerHistory: 30 """A history for banner handling.""" 31 32 _LAST_BANNER_PROMPT_DATE = 'last_banner_prompt_date' 33 34 @staticmethod 35 def create(config_dir: Path) -> BannerHistory: 36 config_dir.mkdir(parents=True, exist_ok=True) 37 history_file = config_dir.joinpath('banner.json') 38 39 if not history_file.exists(): 40 history_file.touch() 41 history = {} 42 else: 43 try: 44 history = json.loads(history_file.read_text()) 45 except json.JSONDecodeError as e: 46 atest_utils.print_and_log_error( 47 'Banner history json file is in a bad format: %s', e 48 ) 49 history = {} 50 51 return BannerHistory(history_file, history) 52 53 def __init__(self, history_file: Path, history: dict): 54 self._history_file = history_file 55 self._history = history 56 57 def get_last_banner_prompt_date(self) -> str: 58 """Get the last date when banner was prompt.""" 59 return self._history.get(BannerHistory._LAST_BANNER_PROMPT_DATE, '') 60 61 def set_last_banner_prompt_date(self, date: str): 62 """Set the last date when banner was prompt.""" 63 self._history[BannerHistory._LAST_BANNER_PROMPT_DATE] = date 64 self._history_file.write_text(json.dumps(self._history)) 65 66 67class BannerPrinter: 68 """A printer used to collect and print banners.""" 69 70 @staticmethod 71 def create() -> BannerPrinter: 72 return BannerPrinter(atest_utils.get_config_folder()) 73 74 def __init__(self, config_dir: Path): 75 self._messages = [] 76 self._config_dir = config_dir 77 78 def register(self, message: str): 79 """Register a banner message.""" 80 self._messages.append(message) 81 82 def print(self, print_func: Callable = None, date_supplier: Callable = None): 83 """Print the banners.""" 84 85 if not self._messages: 86 return 87 88 if not print_func: 89 print_func = lambda m: atest_utils.colorful_print(m, constants.YELLOW) 90 91 if not date_supplier: 92 date_supplier = lambda: str(date.today()) 93 94 today = date_supplier() 95 history = BannerHistory.create(self._config_dir) 96 if history.get_last_banner_prompt_date() != today: 97 for message in self._messages: 98 print_func(message) 99 100 history.set_last_banner_prompt_date(today) 101