More fixes and adptations for pubchem changes

Author: mhucka

src/openfermion/chem/pubchem.py

@@ -9,6 +9,10 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+"""Function to extract geometry from the PubChem database."""
+
+import time
+import urllib.error
 
 
 def geometry_from_pubchem(name: str, structure: str = None):
@@ -32,15 +36,26 @@ def geometry_from_pubchem(name: str, structure: str = None):
     """
     import pubchempy
 
+    def get_compounds_with_retries(name, record_type):
+        max_retries = 3
+        delay = 2
+        for attempt in range(max_retries):
+            try:
+                return pubchempy.get_compounds(name, 'name', record_type=record_type)
+            except (pubchempy.PubChemHTTPError, urllib.error.URLError):
+                if attempt == max_retries - 1:
+                    raise
+                time.sleep(delay)
+
     if structure in ['2d', '3d']:
-        pubchempy_molecule = pubchempy.get_compounds(name, 'name', record_type=structure)
+        pubchempy_molecule = get_compounds_with_retries(name, record_type=structure)
     elif structure is None:
         # Ideally get the 3-D geometry if available.
-        pubchempy_molecule = pubchempy.get_compounds(name, 'name', record_type='3d')
+        pubchempy_molecule = get_compounds_with_retries(name, record_type='3d')
 
         # If the 3-D geometry isn't available, get the 2-D geometry instead.
         if not pubchempy_molecule:
-            pubchempy_molecule = pubchempy.get_compounds(name, 'name', record_type='2d')
+            pubchempy_molecule = get_compounds_with_retries(name, record_type='2d')
     else:
         raise ValueError('Incorrect value for the argument structure=%s' % structure)
 
@@ -52,9 +67,7 @@ def geometry_from_pubchem(name: str, structure: str = None):
         )
         return None
 
-    pubchempy_geometry = pubchempy_molecule[0].to_dict(properties=['atoms'])['atoms']
-    geometry = [
-        (atom['element'], (atom['x'], atom['y'], atom.get('z', 0))) for atom in pubchempy_geometry
-    ]
+    pubchempy_geometry = pubchempy_molecule[0].atoms
+    geometry = [(atom.element, (atom.x, atom.y, atom.z or 0)) for atom in pubchempy_geometry]
 
     return geometry

src/openfermion/chem/pubchem_test.py

@@ -22,12 +22,17 @@
 from openfermion.testing.testing_utils import module_importable
 
 
+class MockAtom:
+    def __init__(self, element, x, y, z=0):
+        self.element = element
+        self.x = x
+        self.y = y
+        self.z = z
+
+
 class MockCompound:
     def __init__(self, atoms):
-        self._atoms = atoms
-
-    def to_dict(self, properties):
-        return {'atoms': self._atoms}
+        self.atoms = [MockAtom(a['element'], a['x'], a['y'], a.get('z', 0)) for a in atoms]
 
 
 def mock_get_compounds(name, searchtype, record_type='2d'):
@@ -70,20 +75,6 @@ def mock_get_compounds(name, searchtype, record_type='2d'):
 
 @using_pubchempy
 class OpenFermionPubChemTest(unittest.TestCase):
-    def _get_geometry_with_retries(self, name):
-        import pubchempy
-        import urllib.error
-
-        max_retries = 3
-        delay = 2
-        for attempt in range(max_retries):
-            try:
-                return geometry_from_pubchem(name)
-            except (pubchempy.PubChemHTTPError, urllib.error.URLError):
-                if attempt == max_retries - 1:
-                    raise
-                time.sleep(delay)
-
     @patch('pubchempy.get_compounds', mock_get_compounds)
     def test_water(self):
         water_geometry = geometry_from_pubchem('water')
@@ -161,7 +152,7 @@ def test_geometry_from_pubchem_live_api(self):
         except ImportError:  # pragma: no cover
             return
 
-        water_geometry = self._get_geometry_with_retries('water')
+        water_geometry = geometry_from_pubchem('water')
         self.assertEqual(len(water_geometry), 3)
 
     def _get_mock_pubchem_error(self):
@@ -181,6 +172,24 @@ def read(self):
 
             return pubchempy.PubChemHTTPError(FakeHTTPError(503, 'Server Busy'))
 
+
+    def _get_mock_pubchem_error(self):
+        import pubchempy
+
+        try:
+            return pubchempy.PubChemHTTPError(503, 'Server Busy', ['busy'])
+        except TypeError:
+
+            class FakeHTTPError(Exception):
+                def __init__(self, code, msg):
+                    self.code = code
+                    self.reason = msg
+
+                def read(self):
+                    return b'{"Fault": {"Details": ["busy"]}}'
+
+            return pubchempy.PubChemHTTPError(FakeHTTPError(503, 'Server Busy'))
+
     @patch('time.sleep', return_value=None)
     @patch('pubchempy.get_compounds')
     def test_geometry_from_pubchem_retry_success(self, mock_get_compounds, mock_sleep):
@@ -218,7 +227,7 @@ def test_geometry_from_pubchem_retry_success(self, mock_get_compounds, mock_slee
             ],
         ]
 
-        water_geometry = self._get_geometry_with_retries('water')
+        water_geometry = geometry_from_pubchem('water')
 
         self.assertEqual(len(water_geometry), 3)
         self.assertEqual(mock_get_compounds.call_count, 3)
@@ -235,7 +244,7 @@ def test_geometry_from_pubchem_retry_failure(self, mock_get_compounds, mock_slee
         mock_get_compounds.side_effect = self._get_mock_pubchem_error()
 
         with self.assertRaises(pubchempy.PubChemHTTPError):
-            self._get_geometry_with_retries('water')
+            geometry_from_pubchem('water')
 
         self.assertEqual(mock_get_compounds.call_count, 3)
         self.assertEqual(mock_sleep.call_count, 2)