Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/et_xmlfile/xmlfile.py: 78%

1from __future__ import absolute_import 

2# Copyright (c) 2010-2015 openpyxl 

3 

4"""Implements the lxml.etree.xmlfile API using the standard library xml.etree""" 

5 

6 

7from contextlib import contextmanager 

8 

9from xml.etree.ElementTree import Element, tostring 

10 

11 

12class LxmlSyntaxError(Exception): 

13 pass 

14 

15 

16class _FakeIncrementalFileWriter(object): 

17 """Replacement for _IncrementalFileWriter of lxml. 

18 Uses ElementTree to build xml in memory.""" 

19 def __init__(self, output_file): 

20 self._element_stack = [] 

21 self._top_element = None 

22 self._file = output_file 

23 self._have_root = False 

24 

25 @contextmanager 

26 def element(self, tag, attrib=None, nsmap=None, **_extra): 

27 """Create a new xml element using a context manager. 

28 The elements are written when the top level context is left. 

29 

30 This is for code compatibility only as it is quite slow. 

31 """ 

32 

33 # __enter__ part 

34 self._have_root = True 

35 if attrib is None: 

36 attrib = {} 

37 self._top_element = Element(tag, attrib=attrib, **_extra) 

38 self._top_element.text = '' 

39 self._top_element.tail = '' 

40 self._element_stack.append(self._top_element) 

41 yield 

42 

43 # __exit__ part 

44 el = self._element_stack.pop() 

45 if self._element_stack: 

46 parent = self._element_stack[-1] 

47 parent.append(self._top_element) 

48 self._top_element = parent 

49 else: 

50 self._write_element(el) 

51 self._top_element = None 

52 

53 def write(self, arg): 

54 """Write a string or subelement.""" 

55 

56 if isinstance(arg, str): 

57 # it is not allowed to write a string outside of an element 

58 if self._top_element is None: 

59 raise LxmlSyntaxError() 

60 

61 if len(self._top_element) == 0: 

62 # element has no children: add string to text 

63 self._top_element.text += arg 

64 else: 

65 # element has children: add string to tail of last child 

66 self._top_element[-1].tail += arg 

67 

68 else: 

69 if self._top_element is not None: 

70 self._top_element.append(arg) 

71 elif not self._have_root: 

72 self._write_element(arg) 

73 else: 

74 raise LxmlSyntaxError() 

75 

76 def _write_element(self, element): 

77 xml = tostring(element) 

78 self._file.write(xml) 

79 

80 def __enter__(self): 

81 pass 

82 

83 def __exit__(self, type, value, traceback): 

84 # without root the xml document is incomplete 

85 if not self._have_root: 

86 raise LxmlSyntaxError() 

87 

88 

89class xmlfile(object): 

90 """Context manager that can replace lxml.etree.xmlfile.""" 

91 def __init__(self, output_file, buffered=False, encoding=None, close=False): 

92 if isinstance(output_file, str): 

93 self._file = open(output_file, 'wb') 

94 self._close = True 

95 else: 

96 self._file = output_file 

97 self._close = close 

98 

99 def __enter__(self): 

100 return _FakeIncrementalFileWriter(self._file) 

101 

102 def __exit__(self, type, value, traceback): 

103 if self._close == True: 

104 self._file.close()