make and test reusable function in pvlib.tools

Author: kandersolar

pvlib/iotools/psm4.py

@@ -11,7 +11,7 @@
 import requests
 import pandas as pd
 from json import JSONDecodeError
-import contextlib
+from pvlib import tools
 
 NSRDB_API_BASE = "https://developer.nrel.gov/api/nsrdb/v2/solar/"
 PSM4_AGG_ENDPOINT = "nsrdb-GOES-aggregated-v4-0-0-download.csv"
@@ -735,14 +735,7 @@ def read_nsrdb_psm4(filename, map_variables=True):
     .. [2] `Standard Time Series Data File Format
        <https://web.archive.org/web/20170207203107/https://sam.nrel.gov/sites/default/files/content/documents/pdf/wfcsv.pdf>`_
     """
-    if hasattr(filename, "read"):
-        # already a file-like object
-        context = contextlib.nullcontext(filename)
-    else:
-        # otherwise, assume a filename or path
-        context = open(str(filename), 'r')
-
-    with context as fbuf:
+    with tools._file_context_manager(filename) as fbuf:
         # The first 2 lines of the response are headers with metadata
         metadata_fields = fbuf.readline().split(',')
         metadata_values = fbuf.readline().split(',')

pvlib/tools.py

@@ -2,6 +2,7 @@
 Collection of functions used in pvlib_python
 """
 
+import contextlib
 import datetime as dt
 import warnings
 
@@ -559,3 +560,29 @@ def normalize_max2one(a):
     except ValueError:  # fails for pandas objects
         res = a.div(a.abs().max(axis=0, skipna=True))
     return res
+
+
+def _file_context_manager(filename_or_object, mode='r'):
+    """
+    Open a filename/path for reading, or pass a file-like object
+    through unchanged.
+
+    Parameters
+    ----------
+    filename_or_object : str, path-like, or file-like object
+        The filename/path or object to convert to an object
+
+    Returns
+    -------
+    context : context manager
+        A file-like object to be used via python's "with [context] as buffer:"
+        syntax.
+    """
+    
+    if hasattr(filename_or_object, "read"):
+        # already a file-like object
+        context = contextlib.nullcontext(filename_or_object)
+    else:
+        # otherwise, assume a filename or path
+        context = open(str(filename_or_object), mode=mode)
+    return context

tests/test_tools.py

@@ -1,4 +1,8 @@
 from datetime import datetime
+from io import StringIO
+import os
+from pathlib import Path
+import tempfile
 from zoneinfo import ZoneInfo
 
 import numpy as np
@@ -252,3 +256,29 @@ def test_djd_to_datetime():
 
     expected = datetime(1974, 6, 22, 23, 30, 15, tzinfo=ZoneInfo("UTC"))
     assert tools.djd_to_datetime(djd) == expected
+
+
+def test__file_context_manager():
+    with tempfile.TemporaryDirectory() as td:
+        # make a test file
+        filename = os.path.join(td, 'test.txt')
+        with open(filename, 'w') as fh:
+            fh.write('test content')
+    
+        # test with filename as string:
+        with tools._file_context_manager(filename) as obj:
+            assert obj.read() == "test content"
+    
+        # test with filename as Path:
+        with tools._file_context_manager(Path(filename)) as obj:
+            assert obj.read() == "test content"
+    
+        # test with file object:
+        with open(filename, "r") as f:
+            with tools._file_context_manager(f) as obj:
+                assert obj.read() == "test content"
+    
+    # test with buffer:
+    buffer = StringIO("test content")
+    with tools._file_context_manager(buffer) as obj:
+        assert obj.read() == "test content"