chendl/compositional_test

1import argparse2import json3import os4import time5import zipfile6 7from get_ci_error_statistics import download_artifact, get_artifacts_links8 9from transformers import logging10 11 12logger = logging.get_logger(__name__)13 14 15def extract_warnings_from_single_artifact(artifact_path, targets):16    """Extract warnings from a downloaded artifact (in .zip format)"""17    selected_warnings = set()18    buffer = []19 20    def parse_line(fp):21        for line in fp:22            if isinstance(line, bytes):23                line = line.decode("UTF-8")24            if "warnings summary (final)" in line:25                continue26            # This means we are outside the body of a warning27            elif not line.startswith(" "):28                # process a single warning and move it to `selected_warnings`.29                if len(buffer) > 0:30                    warning = "\n".join(buffer)31                    # Only keep the warnings specified in `targets`32                    if any(f": {x}: " in warning for x in targets):33                        selected_warnings.add(warning)34                    buffer.clear()35                continue36            else:37                line = line.strip()38                buffer.append(line)39 40    if from_gh:41        for filename in os.listdir(artifact_path):42            file_path = os.path.join(artifact_path, filename)43            if not os.path.isdir(file_path):44                # read the file45                if filename != "warnings.txt":46                    continue47                with open(file_path) as fp:48                    parse_line(fp)49    else:50        try:51            with zipfile.ZipFile(artifact_path) as z:52                for filename in z.namelist():53                    if not os.path.isdir(filename):54                        # read the file55                        if filename != "warnings.txt":56                            continue57                        with z.open(filename) as fp:58                            parse_line(fp)59        except Exception:60            logger.warning(61                f"{artifact_path} is either an invalid zip file or something else wrong. This file is skipped."62            )63 64    return selected_warnings65 66 67def extract_warnings(artifact_dir, targets):68    """Extract warnings from all artifact files"""69 70    selected_warnings = set()71 72    paths = [os.path.join(artifact_dir, p) for p in os.listdir(artifact_dir) if (p.endswith(".zip") or from_gh)]73    for p in paths:74        selected_warnings.update(extract_warnings_from_single_artifact(p, targets))75 76    return selected_warnings77 78 79if __name__ == "__main__":80 81    def list_str(values):82        return values.split(",")83 84    parser = argparse.ArgumentParser()85    # Required parameters86    parser.add_argument("--workflow_run_id", type=str, required=True, help="A GitHub Actions workflow run id.")87    parser.add_argument(88        "--output_dir",89        type=str,90        required=True,91        help="Where to store the downloaded artifacts and other result files.",92    )93    parser.add_argument("--token", default=None, type=str, help="A token that has actions:read permission.")94    # optional parameters95    parser.add_argument(96        "--targets",97        default="DeprecationWarning,UserWarning,FutureWarning",98        type=list_str,99        help="Comma-separated list of target warning(s) which we want to extract.",100    )101    parser.add_argument(102        "--from_gh",103        action="store_true",104        help="If running from a GitHub action workflow and collecting warnings from its artifacts.",105    )106 107    args = parser.parse_args()108 109    from_gh = args.from_gh110    if from_gh:111        # The artifacts have to be downloaded using `actions/download-artifact@v3`112        pass113    else:114        os.makedirs(args.output_dir, exist_ok=True)115 116        # get download links117        artifacts = get_artifacts_links(args.workflow_run_id, token=args.token)118        with open(os.path.join(args.output_dir, "artifacts.json"), "w", encoding="UTF-8") as fp:119            json.dump(artifacts, fp, ensure_ascii=False, indent=4)120 121        # download artifacts122        for idx, (name, url) in enumerate(artifacts.items()):123            print(name)124            print(url)125            print("=" * 80)126            download_artifact(name, url, args.output_dir, args.token)127            # Be gentle to GitHub128            time.sleep(1)129 130    # extract warnings from artifacts131    selected_warnings = extract_warnings(args.output_dir, args.targets)132    selected_warnings = sorted(selected_warnings)133    with open(os.path.join(args.output_dir, "selected_warnings.json"), "w", encoding="UTF-8") as fp:134        json.dump(selected_warnings, fp, ensure_ascii=False, indent=4)135