From 61af4504c37de9b4f45b96e7ae3bf98c062e8b3c Mon Sep 17 00:00:00 2001
From: Jet <38184875+jjshoots@users.noreply.github.com>
Date: Fri, 14 Oct 2022 11:37:24 +0100
Subject: [PATCH] Update setup.py

---
 setup.py | 42 ++++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/setup.py b/setup.py
index 2381e7def..928596ff8 100644
--- a/setup.py
+++ b/setup.py
@@ -4,12 +4,28 @@ import re
 
 from setuptools import find_packages, setup
 
-with open("gymnasium/version.py") as file:
-    full_version = file.read()
-    assert (
-        re.match(r'VERSION = "\d\.\d+\.\d+"\n', full_version).group(0) == full_version
-    ), f"Unexpected version: {full_version}"
-    VERSION = re.search(r"\d\.\d+\.\d+", full_version).group(0)
+def get_description():
+    with open("README.md") as file:
+        long_description = ""
+        header_count = 0
+        for line in file:
+            if line.startswith("##"):
+                header_count += 1
+            if header_count < 2:
+                long_description += line
+            else:
+                break
+    return header_count, long_description
+
+def get_version():
+    """Gets the pettingzoo version."""
+    path = "gymnasium/version.py"
+    with open(path) as file:
+        lines = file.readlines()
+        assert (
+            re.match(r'VERSION = "\d\.\d+\.\d+"\n', full_version).group(0) == full_version
+        ), f"Unexpected version: {full_version}"
+        return re.search(r"\d\.\d+\.\d+", full_version).group(0)
 
 # Environment-specific dependencies.
 extras = {
@@ -35,17 +51,7 @@ extras["all"] = list(
     set(itertools.chain.from_iterable(map(lambda group: extras[group], all_groups)))
 )
 
-# Uses the readme as the description on PyPI
-with open("README.md") as fh:
-    long_description = ""
-    header_count = 0
-    for line in fh:
-        if line.startswith("##"):
-            header_count += 1
-        if header_count < 2:
-            long_description += line
-        else:
-            break
+header_count, long_description = get_description()
 
 setup(
     author="Farama Foundation",
@@ -87,6 +93,6 @@ setup(
     python_requires=">=3.6",
     tests_require=extras["testing"],
     url="https://gymnasium.farama.org/",
-    version=VERSION,
+    version=get_version(),
     zip_safe=False,
 )