fix regression in orgparse.load method for file-like objects

karlicoss · karlicoss · commit 45b366eb5c23 · 2021-01-08T07:10:34.000Z
broken in 3067189 see #32
diff --git a/orgparse/__init__.py b/orgparse/__init__.py
@@ -107,17 +107,19 @@
 # [[[end]]]
 
 import codecs
-from typing import Iterable
+from pathlib import Path
+from typing import Iterable, Union, Optional, TextIO
 
-from .node import parse_lines, OrgNode # todo basenode??
+
+from .node import parse_lines, OrgEnv, OrgNode # todo basenode??
 from .utils.py3compat import basestring
 
 __author__ = 'Takafumi Arakaki, Dmitrii Gerasimov'
 __license__ = 'BSD License'
 __all__ = ["load", "loads", "loadi"]
 
 
-def load(path, env=None):
+def load(path: Union[str, Path, TextIO], env: Optional[OrgEnv]=None) -> OrgNode:
     """
     Load org-mode document from a file.
 
@@ -127,18 +129,18 @@ def load(path, env=None):
     :rtype: :class:`orgparse.node.OrgRootNode`
 
     """
-    path = str(path) # in case of pathlib.Path
-    if isinstance(path, basestring):
-        orgfile = codecs.open(path, encoding='utf8')
-        filename = path
+    orgfile: TextIO
+    if isinstance(path, (str, Path)):
+        orgfile = codecs.open(str(path), encoding='utf8')
+        filename = str(path)
     else:
         orgfile = path
         filename = path.name if hasattr(path, 'name') else '<file-like>'
     return loadi((l.rstrip('\n') for l in orgfile.readlines()),
                  filename=filename, env=env)
 
 
-def loads(string: str, filename='<string>', env=None) -> OrgNode:
+def loads(string: str, filename: str='<string>', env: Optional[OrgEnv]=None) -> OrgNode:
     """
     Load org-mode document from a string.
 
@@ -148,7 +150,7 @@ def loads(string: str, filename='<string>', env=None) -> OrgNode:
     return loadi(string.splitlines(), filename=filename, env=env)
 
 
-def loadi(lines: Iterable[str], filename='<lines>', env=None) -> OrgNode:
+def loadi(lines: Iterable[str], filename: str='<lines>', env: Optional[OrgEnv]=None) -> OrgNode:
     """
     Load org-mode document from an iterative object.
 
diff --git a/orgparse/tests/test_misc.py b/orgparse/tests/test_misc.py
@@ -158,3 +158,14 @@ def test_filetags_are_tags() -> None:
     assert root.tags == {'f1', 'f2'}
     child = root.children[0].children[0]
     assert child.tags == {'f1', 'f2', 'h1'}
+
+
+def test_load_filelike() -> None:
+    import io
+    stream = io.StringIO('''
+* heading1
+* heading 2
+''')
+    root = load(stream)
+    assert len(root.children) == 2
+    assert root.env.filename == '<file-like>'
