From 52b7e5e32554265ec80249985b0dc183a44f37c7 Mon Sep 17 00:00:00 2001
From: Jesse Farebrother <jessefarebro@gmail.com>
Date: Tue, 28 Sep 2021 11:42:26 -0600
Subject: [PATCH] Add extra [accept-rom-license] for AutoROM (#2420)

---
 py.Dockerfile         |  3 +--
 setup.py              | 23 ++++++++++++++---------
 test_requirements.txt |  1 -
 3 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/py.Dockerfile b/py.Dockerfile
index 1a7088e4be1..5e71ef6d647 100644
--- a/py.Dockerfile
+++ b/py.Dockerfile
@@ -15,7 +15,6 @@ ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.mujoco/mjpro150/bin
 COPY . /usr/local/gym/
 WORKDIR /usr/local/gym/
 
-RUN pip install -e .[nomujoco] && pip install -r test_requirements.txt
-RUN AutoROM --accept-license
+RUN pip install -e .[nomujoco,accept-rom-license] && pip install -r test_requirements.txt
 
 ENTRYPOINT ["/usr/local/gym/bin/docker_entrypoint"]
diff --git a/setup.py b/setup.py
index 4e1aa09a3b8..1994b5efcf3 100644
--- a/setup.py
+++ b/setup.py
@@ -1,5 +1,6 @@
 import os.path
 import sys
+import itertools
 
 from setuptools import find_packages, setup
 
@@ -10,6 +11,7 @@
 # Environment-specific dependencies.
 extras = {
     "atari": ["ale-py~=0.7"],
+    "accept-rom-license": ["autorom[accept-rom-license]~=0.4.2"],
     "box2d": ["box2d-py==2.3.5", "pyglet>=1.4.0"],
     "classic_control": ["pyglet>=1.4.0"],
     "mujoco": ["mujoco_py>=1.50, <2.0"],
@@ -19,17 +21,20 @@
 }
 
 # Meta dependency groups.
+nomujoco_blacklist = set(["mujoco", "robotics", "accept-rom-license"])
+nomujoco_groups = set(extras.keys()) - nomujoco_blacklist
+
 extras["nomujoco"] = list(
-    set(
-        [
-            item
-            for name, group in extras.items()
-            if name != "mujoco" and name != "robotics"
-            for item in group
-        ]
-    )
+    itertools.chain.from_iterable(map(lambda group: extras[group], nomujoco_groups))
+)
+
+
+all_blacklist = set(["accept-rom-license"])
+all_groups = set(extras.keys()) - all_blacklist
+
+extras["all"] = list(
+    itertools.chain.from_iterable(map(lambda group: extras[group], all_groups))
 )
-extras["all"] = list(set([item for group in extras.values() for item in group]))
 
 setup(
     name="gym",
diff --git a/test_requirements.txt b/test_requirements.txt
index 81e2b0e1535..7f59ce32d86 100644
--- a/test_requirements.txt
+++ b/test_requirements.txt
@@ -1,4 +1,3 @@
 lz4~=3.1
 pytest~=6.2
 pytest-forked~=1.3
-AutoROM>=0.3