Coverage for src / competitive_verifier / download / download.py: 100%

1from argparse import ArgumentParser 

2from collections.abc import Iterable 

3from itertools import chain 

4from logging import getLogger 

5from typing import Literal 

6 

7from pydantic import Field 

8 

9from competitive_verifier import oj 

10from competitive_verifier.arg import ( 

11 OptionalVerifyFilesJsonArguments, 

12 VerboseArguments, 

13) 

14from competitive_verifier.models import ( 

15 ProblemVerification, 

16 VerificationFile, 

17 VerificationInput, 

18) 

19from competitive_verifier.resource import try_ulimit_stack 

20 

21logger = getLogger(__name__) 

22 

23UrlOrVerificationFile = str | VerificationFile 

24 

25 

26def _parse_single_url(url_or_file: UrlOrVerificationFile) -> Iterable[str]: 

27 if isinstance(url_or_file, str): 

28 yield url_or_file 

29 else: 

30 for v in url_or_file.verification_list: 

31 if isinstance(v, ProblemVerification): 

32 yield v.problem 

33 

34 

35def parse_urls( 

36 url_or_file: UrlOrVerificationFile | Iterable[UrlOrVerificationFile], 

37) -> set[str]: 

38 if isinstance(url_or_file, (str, VerificationFile)): 

39 return set(_parse_single_url(url_or_file)) 

40 

41 return set(chain.from_iterable(map(_parse_single_url, url_or_file))) 

42 

43 

44def download_files( 

45 url_or_file: UrlOrVerificationFile | Iterable[UrlOrVerificationFile], 

46 *, 

47 check: bool = False, 

48 group_log: bool = False, 

49) -> bool: 

50 result = True 

51 try_ulimit_stack() 

52 

53 for url in parse_urls(url_or_file): 

54 if not oj.download(url, group_log=group_log): 

55 result = False 

56 

57 if check and not result: 

58 raise RuntimeError("Failed to download") 

59 return result 

60 

61 

62class Download(OptionalVerifyFilesJsonArguments, VerboseArguments): 

63 subcommand: Literal["download"] = Field( 

64 default="download", 

65 description="Download problems", 

66 ) 

67 urls: list[str] = Field(default_factory=list) 

68 

69 @classmethod 

70 def add_parser(cls, parser: ArgumentParser): 

71 super().add_parser(parser) 

72 parser.add_argument( 

73 "urls", 

74 nargs="*", 

75 help="A list of problem URL", 

76 ) 

77 

78 def _run(self) -> bool: 

79 logger.debug("arguments:%s", self) 

80 logger.info( 

81 "verify_files_json=%s, urls=%s", 

82 str(self.verify_files_json), 

83 self.urls, 

84 ) 

85 files: list[VerificationFile] = [] 

86 if self.verify_files_json: 

87 files.extend( 

88 VerificationInput.parse_file_relative( 

89 self.verify_files_json 

90 ).files.values() 

91 ) 

92 

93 return download_files(files + self.urls, group_log=True)