Coverage for app/utils.py: 64%

1import sys 

2from datetime import datetime, timedelta 

3from typing import Literal, cast 

4from urllib.parse import parse_qs, urlencode, urlparse 

5 

6from pydantic import ValidationError 

7from sqlalchemy import text 

8 

9from app.core.db import read_engine 

10from app.core.logger import logger 

11from app.models import QueryParams, TimeRange, TimeRangeValidValues 

12 

13 

14def start_of_week(date: str | datetime) -> str: 

15 if isinstance(date, str): 

16 date = datetime.strptime(date, "%Y-%m-%d") 

17 start = date - timedelta(days=date.weekday()) 

18 return start.strftime("%Y-%m-%d") 

19 

20 

21def validate_package(package_name: str) -> bool: 

22 with read_engine.connect() as conn: 

23 result = conn.execute( 

24 text(""" 

25 select 

26 package_name 

27 from 

28 pypi_packages 

29 where 

30 package_name = :package_name 

31 """), 

32 {"package_name": package_name}, 

33 ) 

34 exists = result.scalar() is not None 

35 if not exists: 

36 logger.info(f"Package validation failed for: {package_name}") 

37 return exists 

38 

39 

40def parse_query_params(url: str) -> QueryParams: 

41 try: 

42 qs = parse_qs(urlparse(url).query) 

43 packages = qs.get("packages", []) 

44 time_range = TimeRange( 

45 value=cast(TimeRangeValidValues, qs.get("time_range", ["3months"])[0]) 

46 ) 

47 show_percentage = cast(Literal["on"], qs.get("show_percentage", [None])[0]) 

48 return QueryParams( 

49 time_range=time_range, packages=packages, show_percentage=show_percentage 

50 ) 

51 except ValidationError as e: 

52 return QueryParams(error=str(e)) 

53 

54 

55def generate_hx_push_url(query_params: QueryParams) -> str: 

56 if not query_params.packages: 

57 return "/" 

58 params = { 

59 "packages": [package for package in query_params.packages], 

60 "time_range": query_params.time_range.value, 

61 } 

62 if query_params.show_percentage is not None: 

63 params["show_percentage"] = [query_params.show_percentage] 

64 return "?" + urlencode(params, doseq=True) 

65 

66 

67def extract_last_script_tag(html_content: str) -> str | None: 

68 pos = len(html_content) 

69 while True: 

70 pos = html_content.rfind("<script", 0, pos) 

71 if pos == -1: 

72 return None 

73 

74 end_pos = html_content.find("</script>", pos) 

75 if end_pos != -1: 

76 return html_content[pos : end_pos + len("</script>")] 

77 

78 

79def validate_date(date_text: str) -> bool: 

80 try: 

81 datetime.strptime(date_text, "%Y-%m-%d") 

82 return True 

83 except ValueError: 

84 return False 

85 

86 

87def parse_dates(dates: tuple[str, str]) -> tuple[str, str]: 

88 start_date, end_date = dates[0], dates[1] 

89 

90 if not all(validate_date(date) for date in (start_date, end_date)): 

91 logger.error("Start date and end date must be in YYYY-MM-DD format") 

92 sys.exit(1) 

93 

94 if end_date < start_date: 

95 logger.error("End date must be greater than or equal to the start date") 

96 sys.exit(1) 

97 

98 return start_date, end_date