Coverage for polar/kit/csv.py: 41%

1import collections 

2import csv 

3from collections.abc import Iterable 

4from typing import TYPE_CHECKING, Any, BinaryIO 

5 

6if TYPE_CHECKING: 6 ↛ 7line 6 didn't jump to line 7 because the condition on line 6 was never true

7 import _csv 

8 

9from .email import EmailNotValidError, validate_email 

10 

11 

12def get_iterable_from_binary_io(file: BinaryIO) -> Iterable[str]: 

13 for line in file: 

14 yield line.decode("utf-8") 

15 

16 

17def get_emails_from_csv(lines: Iterable[str]) -> set[str]: 

18 emails: set[str] = set() 

19 

20 reader = csv.DictReader(lines) 

21 if reader.fieldnames is None: 

22 return emails 

23 

24 try: 

25 email_field = next( 

26 field for field in reader.fieldnames if "email" in field.lower() 

27 ) 

28 except StopIteration: 

29 return emails 

30 

31 for row in reader: 

32 email = row.get(email_field) 

33 if email is not None: 

34 try: 

35 validate_email(email) 

36 except EmailNotValidError: 

37 continue 

38 else: 

39 emails.add(email) 

40 

41 return emails 

42 

43 

44class IterableCSVWriter: 

45 """ 

46 Utility class wrapping the built-in csv.writer allowing 

47 to generate CSV rows as strings. 

48 

49 It's useful to generate CSV with StreamingResponse, for example. 

50 """ 

51 

52 writer: "_csv._writer" 

53 

54 def __init__( 

55 self, 

56 dialect: "_csv._DialectLike" = "excel", 

57 *, 

58 delimiter: str = ",", 

59 quotechar: str | None = '"', 

60 escapechar: str | None = None, 

61 doublequote: bool = True, 

62 skipinitialspace: bool = False, 

63 lineterminator: str = "\r\n", 

64 quoting: "_csv._QuotingType" = 0, 

65 strict: bool = False, 

66 ) -> None: 

67 self._lines: collections.deque[str] = collections.deque() 

68 self.writer = csv.writer( 

69 self, 

70 dialect=dialect, 

71 delimiter=delimiter, 

72 quotechar=quotechar, 

73 escapechar=escapechar, 

74 doublequote=doublequote, 

75 skipinitialspace=skipinitialspace, 

76 lineterminator=lineterminator, 

77 quoting=quoting, 

78 strict=strict, 

79 ) 

80 

81 def getrow(self, row: Iterable[Any]) -> str: 

82 self.writer.writerow(row) 

83 return self.read() 

84 

85 def write(self, line: str) -> None: 

86 self._lines.append(line) 

87 

88 def read(self) -> str: 

89 return self._lines.popleft()