├── .gitignore ├── MANIFEST.in ├── screenshots ├── rtyaml.png └── without-rtyaml.png ├── CONTRIBUTING.md ├── setup.py ├── test.yaml ├── README.md ├── LICENSE └── rtyaml └── __init__.py /.gitignore: -------------------------------------------------------------------------------- 1 | build/ 2 | dist/ 3 | rtyaml.egg-info/ 4 | __pycache__ 5 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include rtyaml/__init__.py 2 | include LICENSE README.rst 3 | -------------------------------------------------------------------------------- /screenshots/rtyaml.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unitedstates/rtyaml/HEAD/screenshots/rtyaml.png -------------------------------------------------------------------------------- /screenshots/without-rtyaml.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/unitedstates/rtyaml/HEAD/screenshots/without-rtyaml.png -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | ## Public domain 2 | 3 | The project is in the public domain within the United States, and copyright and related rights in the work worldwide are waived through the [CC0 1.0 Universal public domain dedication][CC0]. 4 | 5 | All contributions to this project must be released under the CC0 dedication. By submitting a pull request, you are agreeing to comply with this waiver of copyright interest. 6 | 7 | [CC0]: http://creativecommons.org/publicdomain/zero/1.0/ 8 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | # python3 -m pip install --user --upgrade setuptools wheel twine 4 | # rm -rf dist && python3 setup.py sdist bdist_wheel --universal 5 | # twine upload dist/* 6 | # git tag v1.0.XXX 7 | # git push --tags 8 | 9 | from setuptools import setup 10 | 11 | setup( 12 | name='rtyaml', 13 | version='1.0.0', 14 | author=u'Joshua Tauberer', 15 | author_email=u'jt@occams.info', 16 | packages=['rtyaml'], 17 | url='https://github.com/unitedstates/rtyaml', 18 | license='CC0 (copyright waived)', 19 | description='All the annoying things to make YAML usable in a source controlled environment.', 20 | long_description=open("README.md").read(), 21 | long_description_content_type="text/markdown", 22 | install_requires=["pyyaml"], 23 | ) 24 | -------------------------------------------------------------------------------- /test.yaml: -------------------------------------------------------------------------------- 1 | # Comments at the start of the file are preserved 2 | # when round-tripping files. 3 | keys: 4 | zz: 1 # <--- Order of keys is preserved. 5 | yy: 2 6 | xx: 3 7 | strings: 8 | a: "07" # <--- automatically quoted when it looks like an integer 9 | b: "08" # <--- quotes used consistently 10 | c: No quotes are needed when the value can only be a string. 11 | nulls: ~ # <--- tilde 12 | multiline-strings: 13 | short: | 14 | Multi-line text with short lines is 15 | dumped using the "literal" ("|"") block style, 16 | because wrapping isn't needed. 17 | long: > 18 | Text with very long lines automatically uses the folded (">") block style 19 | which wraps long lines automatically so the file fits nicely in a text editor. 20 | 21 | This is a standard YAML format, but you need rtyaml to automatically use it. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | rtyaml: Round-trippable YAML 2 | ============================ 3 | 4 | Primary author: Joshua Tauberer 5 | 6 | - GitHub: 7 | - PyPi: 8 | 9 | This module is a wrapper around `pyyaml` to read and write YAML files with some improvements: 10 | 11 | - Round-tripping YAML files is possible by preserving the order of keys. In Python 3.7+, just use regular `dict`s. In prior versions of Python, use `collections.OrderedDict`. 12 | - Saner defaults are set for how strings are written out to YAML (see below). 13 | - Other sane defaults are chosen like using the "safe" loader/dumper. 14 | - A comment block found at the very beginning of a stream when loading YAML is preserved when writing it back out. 15 | 16 | ## What do you mean by round-tripping? 17 | 18 | Round-tripping is when you load a file and then save it unchanged, you expect the bytes on disk not to change. This isn't possible with PyYAML, and it makes it difficult to use YAML files with version control like git because every time you load and save the file, things can get rearranged. Keys can change order, string quoting styles can change, #-comments are removed, and so on. 19 | 20 | Although `rtyaml` can't provide round-tripping for all files, it does set some sane defaults on PyYAML so that it's easier to achieve. For instance, if you load this file with PyYAML: 21 | 22 | ![](screenshots/rtyaml.png) 23 | 24 | and then save it back out unchanged: 25 | 26 | ```python 27 | import yaml 28 | print(yaml.dump(yaml.load(open('example.yaml')))) 29 | ``` 30 | 31 | you get this mess: 32 | 33 | ![](screenshots/without-rtyaml.png) 34 | 35 | Notice how the comment is gone, the keys `zz`, `yy`, `xx` changed order, the strings are inconsistently formatted, nulls use a confusing keyword, and mappings are condensed into single lines. 36 | 37 | With `ryaml`, you actually get the original file back! That's basically the whole point of this library. 38 | 39 | 40 | ## Installation and usage 41 | 42 | Install: 43 | 44 | pip install rtyaml 45 | 46 | Usage: 47 | 48 | import rtyaml 49 | 50 | with open("myfile.yaml") as f: 51 | stuff = rtyaml.load(f) 52 | 53 | # ...do things to stuf... 54 | 55 | with open("myfile.yaml", "w") as f: 56 | rtyaml.dump(stuff, f) 57 | 58 | As in the underlying pyyaml library, `load` accepts a string or bytes-string containing YAML or an open file object (binary or text). Also, the second argument to `dump` is optional and if omitted the function returns the YAML in a string. 59 | 60 | `load_all` and `dump_all` are also supported, which load and save lists of documents using YAML's `---` document separator. 61 | 62 | Dependencies 63 | ------------ 64 | 65 | - pyyaml (in Ubuntu, the `python-yaml` or `python3-yaml` package) 66 | - libyaml (in Ubuntu, the `libyaml-0-2` package plus, at install time only, `libyaml-dev`) 67 | 68 | Details 69 | ------- 70 | 71 | This library does the following: 72 | 73 | - Uses the native libyaml CSafeLoader and CDumper for both speed and trustable operations. 74 | - Preserves the order of keys in `dicts` rather than alphebetizing the keys (Python >=3.7). 75 | - Allows you to use `collections.OrderedDict`s with `dump` to preserve key order (useful before Python 3.7). 76 | - Writes multi-line strings in block mode (rather than quoted with ugly escaped newline characters), choosing between the literal or folded mode depending on what looks better for the length of the lines in the string. 77 | - Writes mappings and lists in the expanded (one per line) format, which is nice when the output is going in version control. 78 | - Modifies the flow string quoting rules so that any string made up of digits is serialized with quotes. (The default settings serialize the string "01" with quotes but the string "09" without quotes! (Can you figure out why?)) 79 | - `None` is serialized as the tilde rather than as `null`, which is less confusing. 80 | - If a block comment appears at the start of the file (i.e. one or more lines starting with a '#'), write it back out if the same object is written with rtyaml.dump(). 81 | 82 | For Python 3.6 and earlier: 83 | 84 | - Loads mappings `collections.OrderedDict` so that the key order remains the same when dumping the file later using. (This is no longer needed in Python 3.7 because key order is preserved in regular `dict`s now.) 85 | 86 | With-Block Helper for Editing Files In-Place 87 | -------------------------------------------- 88 | 89 | The `rtyaml.edit` class is a utility class that can be used with with blocks that makes it easier to edit YAML files in-place. For example: 90 | 91 | ``` 92 | with rtyaml.edit("path/to/data.yaml", default={}) as data: 93 | data["hello"] = "world" 94 | ``` 95 | 96 | The file is opened for editing ("r+" mode, or "w+" mode if it doesn't exist and a default value is given) and its contents is parsed and returned as the data with-block variable. The file is kept open while the with-block is executing. When the with-block exits, the with-block variable is written back to the file as YAML, and then the file is closed. 97 | 98 | This will, of course, only work if the file contains an array or object (dict), and you cannot assign a new value to the with-block variable (that's just how Python with blocks work). You can only call its methods, i.e., you can edit the list (append, pop, sort, etc.) and dict (get/set keys), but you can't replace the value with an entirely new list or dict. 99 | 100 | If the default parameter is not given, or is None, the file must exist. Otherwise, if the file doesn't exist, it's created and the with-block variable will start you off with the default value. 101 | 102 | You can also pass a stream as the first argument if you want to open the file yourself. The stream must support seek, truncate, and close. If you open a file, you should use the "r+" or "w+" mode. 103 | 104 | Public domain dedication 105 | ------------------------ 106 | 107 | This project is dedicated to the public domain, as indicated in the LICENSE file: 108 | 109 | > The project is in the public domain within the United States, and copyright and related rights in the work worldwide are waived through the CC0 1.0 Universal public domain dedication. 110 | 111 | All contributions to this project must be released under the CC0 dedication. By submitting a pull request, you are agreeing to comply with this waiver of copyright interest. 112 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Creative Commons Legal Code 2 | 3 | CC0 1.0 Universal 4 | 5 | CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE 6 | LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN 7 | ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS 8 | INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES 9 | REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS 10 | PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM 11 | THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED 12 | HEREUNDER. 13 | 14 | Statement of Purpose 15 | 16 | The laws of most jurisdictions throughout the world automatically confer 17 | exclusive Copyright and Related Rights (defined below) upon the creator 18 | and subsequent owner(s) (each and all, an "owner") of an original work of 19 | authorship and/or a database (each, a "Work"). 20 | 21 | Certain owners wish to permanently relinquish those rights to a Work for 22 | the purpose of contributing to a commons of creative, cultural and 23 | scientific works ("Commons") that the public can reliably and without fear 24 | of later claims of infringement build upon, modify, incorporate in other 25 | works, reuse and redistribute as freely as possible in any form whatsoever 26 | and for any purposes, including without limitation commercial purposes. 27 | These owners may contribute to the Commons to promote the ideal of a free 28 | culture and the further production of creative, cultural and scientific 29 | works, or to gain reputation or greater distribution for their Work in 30 | part through the use and efforts of others. 31 | 32 | For these and/or other purposes and motivations, and without any 33 | expectation of additional consideration or compensation, the person 34 | associating CC0 with a Work (the "Affirmer"), to the extent that he or she 35 | is an owner of Copyright and Related Rights in the Work, voluntarily 36 | elects to apply CC0 to the Work and publicly distribute the Work under its 37 | terms, with knowledge of his or her Copyright and Related Rights in the 38 | Work and the meaning and intended legal effect of CC0 on those rights. 39 | 40 | 1. Copyright and Related Rights. A Work made available under CC0 may be 41 | protected by copyright and related or neighboring rights ("Copyright and 42 | Related Rights"). Copyright and Related Rights include, but are not 43 | limited to, the following: 44 | 45 | i. the right to reproduce, adapt, distribute, perform, display, 46 | communicate, and translate a Work; 47 | ii. moral rights retained by the original author(s) and/or performer(s); 48 | iii. publicity and privacy rights pertaining to a person's image or 49 | likeness depicted in a Work; 50 | iv. rights protecting against unfair competition in regards to a Work, 51 | subject to the limitations in paragraph 4(a), below; 52 | v. rights protecting the extraction, dissemination, use and reuse of data 53 | in a Work; 54 | vi. database rights (such as those arising under Directive 96/9/EC of the 55 | European Parliament and of the Council of 11 March 1996 on the legal 56 | protection of databases, and under any national implementation 57 | thereof, including any amended or successor version of such 58 | directive); and 59 | vii. other similar, equivalent or corresponding rights throughout the 60 | world based on applicable law or treaty, and any national 61 | implementations thereof. 62 | 63 | 2. Waiver. To the greatest extent permitted by, but not in contravention 64 | of, applicable law, Affirmer hereby overtly, fully, permanently, 65 | irrevocably and unconditionally waives, abandons, and surrenders all of 66 | Affirmer's Copyright and Related Rights and associated claims and causes 67 | of action, whether now known or unknown (including existing as well as 68 | future claims and causes of action), in the Work (i) in all territories 69 | worldwide, (ii) for the maximum duration provided by applicable law or 70 | treaty (including future time extensions), (iii) in any current or future 71 | medium and for any number of copies, and (iv) for any purpose whatsoever, 72 | including without limitation commercial, advertising or promotional 73 | purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each 74 | member of the public at large and to the detriment of Affirmer's heirs and 75 | successors, fully intending that such Waiver shall not be subject to 76 | revocation, rescission, cancellation, termination, or any other legal or 77 | equitable action to disrupt the quiet enjoyment of the Work by the public 78 | as contemplated by Affirmer's express Statement of Purpose. 79 | 80 | 3. Public License Fallback. Should any part of the Waiver for any reason 81 | be judged legally invalid or ineffective under applicable law, then the 82 | Waiver shall be preserved to the maximum extent permitted taking into 83 | account Affirmer's express Statement of Purpose. In addition, to the 84 | extent the Waiver is so judged Affirmer hereby grants to each affected 85 | person a royalty-free, non transferable, non sublicensable, non exclusive, 86 | irrevocable and unconditional license to exercise Affirmer's Copyright and 87 | Related Rights in the Work (i) in all territories worldwide, (ii) for the 88 | maximum duration provided by applicable law or treaty (including future 89 | time extensions), (iii) in any current or future medium and for any number 90 | of copies, and (iv) for any purpose whatsoever, including without 91 | limitation commercial, advertising or promotional purposes (the 92 | "License"). The License shall be deemed effective as of the date CC0 was 93 | applied by Affirmer to the Work. Should any part of the License for any 94 | reason be judged legally invalid or ineffective under applicable law, such 95 | partial invalidity or ineffectiveness shall not invalidate the remainder 96 | of the License, and in such case Affirmer hereby affirms that he or she 97 | will not (i) exercise any of his or her remaining Copyright and Related 98 | Rights in the Work or (ii) assert any associated claims and causes of 99 | action with respect to the Work, in either case contrary to Affirmer's 100 | express Statement of Purpose. 101 | 102 | 4. Limitations and Disclaimers. 103 | 104 | a. No trademark or patent rights held by Affirmer are waived, abandoned, 105 | surrendered, licensed or otherwise affected by this document. 106 | b. Affirmer offers the Work as-is and makes no representations or 107 | warranties of any kind concerning the Work, express, implied, 108 | statutory or otherwise, including without limitation warranties of 109 | title, merchantability, fitness for a particular purpose, non 110 | infringement, or the absence of latent or other defects, accuracy, or 111 | the present or absence of errors, whether or not discoverable, all to 112 | the greatest extent permissible under applicable law. 113 | c. Affirmer disclaims responsibility for clearing rights of other persons 114 | that may apply to the Work or any use thereof, including without 115 | limitation any person's Copyright and Related Rights in the Work. 116 | Further, Affirmer disclaims responsibility for obtaining any necessary 117 | consents, permissions or other rights required for any use of the 118 | Work. 119 | d. Affirmer understands and acknowledges that Creative Commons is not a 120 | party to this document and has no duty or obligation with respect to 121 | this CC0 or use of the Work. 122 | -------------------------------------------------------------------------------- /rtyaml/__init__.py: -------------------------------------------------------------------------------- 1 | import sys, re, io 2 | 3 | # Import YAML's SafeLoader. 4 | import yaml 5 | try: 6 | # Use the native code backends, if available. 7 | from yaml import CSafeLoader as Loader, CDumper as Dumper 8 | except ImportError: 9 | from yaml import SafeLoader as Loader, Dumper 10 | 11 | # Before Python 3.7 when dicts began to preserver the insertion 12 | # order of elements, in order to preserve the order of mappings, 13 | # YAML must be hooked to load mappings as OrderedDicts. Adapted from: 14 | # https://gist.github.com/317164 15 | from collections import OrderedDict 16 | if sys.version_info < (3,7): 17 | def construct_odict(load, node): 18 | omap = OrderedDict() 19 | yield omap 20 | if not isinstance(node, yaml.MappingNode): 21 | raise yaml.constructor.ConstructorError( 22 | "while constructing an ordered map", 23 | node.start_mark, 24 | "expected a map, but found %s" % node.id, node.start_mark 25 | ) 26 | for key, value in node.value: 27 | key = load.construct_object(key) 28 | value = load.construct_object(value) 29 | omap[key] = value 30 | 31 | Loader.add_constructor(u'tag:yaml.org,2002:map', construct_odict) 32 | 33 | else: 34 | # Starting with Pyhon 3.7, dicts preserve order. But PyYAML by default 35 | # sorts the keys of any mapping it gets when that mapping has an 'items' 36 | # attribute. We override that by adding an explicit representer for 'dict' 37 | # that passes the items directly (so that the value it seems does not have 38 | # an 'items' attribute itself). See https://github.com/yaml/pyyaml/blob/e471e86bf6dabdad45a1438c20a4a5c033eb9034/lib/yaml/representer.py#L119. 39 | # See below for similar code for OrderedDicts. 40 | Dumper.add_representer(dict, lambda self, data: self.represent_mapping('tag:yaml.org,2002:map', data.items())) 41 | 42 | # Tell YAML to serialize OrderedDicts as mappings. Prior to Python 3.7, this 43 | # was the data type that must be used to specify key order, and, per the block 44 | # above, we would deserialize to OrderedDicts. In Python >=3.7, we don't deserialize 45 | # to OrderedDicts anymore but we still allow OrderedDicts to be used in data passed 46 | # to dump. See the block above for why we also add an explicit serializer for dicts. 47 | def ordered_dict_serializer(self, data): 48 | return self.represent_mapping('tag:yaml.org,2002:map', data.items()) 49 | Dumper.add_representer(OrderedDict, ordered_dict_serializer) 50 | 51 | # The standard PyYAML representer for strings does something weird: 52 | # If a value cannot be parsed as another data type, quotes are omitted. 53 | # 54 | # This is incredibly odd when the value is integer-like with a leading 55 | # zero. These values are typically parsed as octal integers, meaning 56 | # quotes would normally be required (that's good). But when the value 57 | # has an '8' or '9' in it, this would make it an invalid octal number 58 | # and so quotes would no longer be required (that's confusing). 59 | # We will override str and unicode output to choose the quotation 60 | # style with our own logic. (According to PyYAML, style can be one of 61 | # the empty string, ', ", |, or >, or None to, presumably, choose 62 | # automatically and use no quoting where possible.) 63 | def our_string_representer(dumper, value): 64 | # let PyYAML choose by default, using no quoting where possible 65 | style = None 66 | 67 | # If it looks like an octal number, force '-quote style. 68 | if re.match(r"^0\d*$", value): style = "'" 69 | 70 | # If it has newlines, request a block style. 71 | if "\n" in value: 72 | # If the average length of a line is very long, then use the folded 73 | # style so that in our output the lines get folded. The drawback when 74 | # this is used on shortlines is that newlines get doubled. So when 75 | # the lines are short, use the literal block style. 76 | lines = value.split("\n") 77 | avg_line_length = sum(len(line) for line in lines) / float(len(lines)) 78 | if avg_line_length > 70: 79 | style = ">" # folded 80 | else: 81 | style = "|" 82 | 83 | return dumper.represent_scalar(u'tag:yaml.org,2002:str', value, style=style) 84 | Dumper.add_representer(str, our_string_representer) 85 | 86 | # Add a representer for nulls too. YAML accepts "~" for None, but the 87 | # default output converts that to "null". Override to always use "~". 88 | Dumper.add_representer(type(None), lambda dumper, value : \ 89 | dumper.represent_scalar(u'tag:yaml.org,2002:null', u"~")) 90 | 91 | # Use a subclss of list when trying to hold onto a block comment at the 92 | # start of a stream. Make sure it serializes back to a plain YAML list. 93 | class RtYamlList(list): 94 | pass 95 | def RtYamlList_serializer(self, data): 96 | return self.represent_sequence('tag:yaml.org,2002:seq', data) 97 | Dumper.add_representer(RtYamlList, RtYamlList_serializer) 98 | 99 | # Provide some wrapper methods that apply typical settings. 100 | 101 | def load(stream): 102 | return do_load(stream, yaml.load) 103 | 104 | def load_all(stream): 105 | return do_load(stream, yaml.load_all) 106 | 107 | def do_load(stream, load_func): 108 | # Read any comment block at the start. 109 | initial_comment_block = "" 110 | 111 | # Try reading the stream until the first non-comment line, then seek back to the 112 | # start of that line. 113 | try: 114 | while True: 115 | p = stream.tell() 116 | line = stream.readline() 117 | if not line or line[0] != "#": 118 | stream.seek(p) 119 | break 120 | initial_comment_block += line 121 | except (AttributeError, io.UnsupportedOperation) as e: 122 | # seek, tell, and readline may not be present and will raise an AttributeError 123 | # or seek/tell may raise io.UnsupportedOperation. 124 | pass 125 | 126 | # Read the object from the stream. 127 | obj = load_func(stream, Loader=Loader) 128 | 129 | # Attach our initial comment to the object so we can save it later (assuming 130 | # this object is written back out). 131 | if initial_comment_block: 132 | if isinstance(obj, list): 133 | # The list class can't be assigned any custom attributes, but we can make a subclass that can. 134 | # Clone the list object into a RtYamlList instance. 135 | obj = RtYamlList(obj) 136 | if isinstance(obj, dict): 137 | # The dict class can't be assigned any custom attributes, so we'll use an OrderedDict instead, which can. 138 | obj = OrderedDict(obj) 139 | obj.__initial_comment_block = initial_comment_block 140 | 141 | return obj 142 | 143 | def dump(data, stream=None): 144 | return do_dump(data, stream, yaml.dump) 145 | 146 | def dump_all(data, stream=None): 147 | return do_dump(data, stream, yaml.dump_all) 148 | 149 | def do_dump(data, stream, dump_func): 150 | # If we pulled in an initial comment block when reading the stream, write 151 | # it back out at the start of the stream. If we're dumping to a string, 152 | # then stream is none. 153 | if hasattr(data, '__initial_comment_block') and stream is not None: 154 | stream.write(data.__initial_comment_block) 155 | 156 | # Write the object to the stream. 157 | ret = dump_func(data, stream, default_flow_style=False, allow_unicode=True, Dumper=Dumper) 158 | 159 | # If we're dumping to a stream, pre-pend the initial comment block. 160 | if hasattr(data, '__initial_comment_block') and stream is None and isinstance(ret, str): 161 | ret = data.__initial_comment_block + ret 162 | 163 | return ret 164 | 165 | def pprint(data): 166 | yaml.dump(data, sys.stdout, default_flow_style=False, allow_unicode=True) 167 | 168 | # This class is a helper that facilitates editing YAML files 169 | # in place. Use as: 170 | # 171 | # with rtyaml.edit("path/to/data.yaml", default={}) as data: 172 | # data["hello"] = "world" 173 | # 174 | # The file is opened for editing ("r+ mode"), read, and its 175 | # parsed YAML data returned as the with-block variable (`data` 176 | # in this case). The file is kept open while the with-block 177 | # is executing. When the with-block exits the stream is overwritten 178 | # with the new YAML data, and the file is closed. 179 | # 180 | # This will of course only work if the file contains an array 181 | # or object (dict), and you cannot assign a new value to the 182 | # with-block variable --- you can only call its methods, e.g. 183 | # you can edit the list and dict, but you can't replace the 184 | # value with an entirely new list or dict. 185 | # 186 | # If the default parameter is not given, or is None, the file 187 | # must exist. Otherwise, if the file doesn't exist, it's created 188 | # and the with-block variable holds the default value. 189 | # 190 | # You can also pass a stream as the first argument if you want 191 | # to open the file yourself. The stream must support seek, truncate, 192 | # and close. 193 | class edit: 194 | def __init__(self, fn_or_stream, default=None): 195 | self.fn_or_stream = fn_or_stream 196 | self.default = default 197 | def __enter__(self): 198 | if isinstance(self.fn_or_stream, str): 199 | # Open the named file. 200 | try: 201 | self.stream = open(self.fn_or_stream, "r+") 202 | except FileNotFoundError: 203 | if not isinstance(self.default, (list, dict)): 204 | # If there is no default and the file 205 | # does not exist, re-raise the exception. 206 | raise 207 | else: 208 | # Create a new file holding the default, 209 | # then seek back to the beginning so 210 | # we can read it below. 211 | self.stream = open(self.fn_or_stream, "w+") 212 | dump(self.default, self.stream) 213 | self.stream.seek(0) 214 | 215 | self.close_on_exit = True 216 | else: 217 | # Use the given stream. 218 | self.stream = self.fn_or_stream 219 | 220 | # Parse stream and return data. 221 | self.data = load(self.stream) 222 | return self.data 223 | 224 | def __exit__(self, *exception): 225 | # Truncate stream and write new data. 226 | self.stream.seek(0); 227 | self.stream.truncate() 228 | dump(self.data, self.stream) 229 | 230 | # Close stream if we opened it. 231 | if getattr(self, "close_on_exit", False): 232 | self.stream.close() 233 | --------------------------------------------------------------------------------