Add flag to get_nwb_version to include prerelease info (#1639)

src/pynwb/io/utils.py

@@ -4,14 +4,19 @@
 from hdmf.build import Builder
 
 
-def get_nwb_version(builder: Builder) -> Tuple[int, ...]:
+def get_nwb_version(builder: Builder, include_prerelease=False) -> Tuple[int, ...]:
     """Get the version of the NWB file from the root of the given builder, as a tuple.
 
     If the "nwb_version" attribute on the root builder equals "2.5.1", then (2, 5, 1) is returned.
-    If the "nwb_version" attribute on the root builder equals "2.5.1-alpha", then (2, 5, 1) is returned.
+    If the "nwb_version" attribute on the root builder equals "2.5.1-alpha" and include_prerelease=False,
+    then (2, 5, 1) is returned.
+    If the "nwb_version" attribute on the root builder equals "2.5.1-alpha" and include_prerelease=True,
+    then (2, 5, 1, "alpha") is returned.
 
     :param builder: Any builder within an NWB file.
     :type builder: Builder
+    :param include_prerelease: Whether to include prerelease information in the returned tuple.
+    :type include_prerelease: bool
     :return: The version of the NWB file, as a tuple.
     :rtype: tuple
     :raises ValueError: if the 'nwb_version' attribute is missing from the root of the NWB file.
@@ -23,5 +28,9 @@ def get_nwb_version(builder: Builder) -> Tuple[int, ...]:
     nwb_version = root_builder.attributes.get("nwb_version")
     if nwb_version is None:
         raise ValueError("'nwb_version' attribute is missing from the root of the NWB file.")
-    nwb_version = re.match(r"(\d+\.\d+\.\d+)", nwb_version)[0]  # trim off any non-numeric symbols at end
-    return tuple([int(i) for i in nwb_version.split(".")])
+    nwb_version_match = re.match(r"(\d+\.\d+\.\d+)", nwb_version)[0]  # trim off any non-numeric symbols at end
+    version_list = [int(i) for i in nwb_version_match.split(".")]
+    if include_prerelease:
+        prerelease_info = nwb_version[nwb_version.index("-")+1:]
+        version_list.append(prerelease_info)
+    return tuple(version_list)

tests/integration/utils/test_io_utils.py

@@ -28,3 +28,21 @@ def test_get_nwb_version_missing(self):
 
         with pytest.raises(ValueError, match="'nwb_version' attribute is missing from the root of the NWB file."):
             get_nwb_version(builder1)
+
+    def test_get_nwb_version_prerelease_false(self):
+        """Get the NWB version from a builder."""
+        builder1 = GroupBuilder(name="root")
+        builder1.set_attribute(name="nwb_version", value="2.0.0-alpha")
+        assert get_nwb_version(builder1) == (2, 0, 0)
+
+    def test_get_nwb_version_prerelease_true1(self):
+        """Get the NWB version from a builder."""
+        builder1 = GroupBuilder(name="root")
+        builder1.set_attribute(name="nwb_version", value="2.0.0-alpha")
+        assert get_nwb_version(builder1, include_prerelease=True) == (2, 0, 0, "alpha")
+
+    def test_get_nwb_version_prerelease_true2(self):
+        """Get the NWB version from a builder."""
+        builder1 = GroupBuilder(name="root")
+        builder1.set_attribute(name="nwb_version", value="2.0.0-alpha.sha-test.5114f85")
+        assert get_nwb_version(builder1, include_prerelease=True) == (2, 0, 0, "alpha.sha-test.5114f85")