with_probabilities support

Author: josiahkhor

setup.py

@@ -8,7 +8,7 @@ def readme():
         return f.read()
 
 setuptools.setup(name='unmand',
-    version='2.1.0',
+    version='2.1.1',
     description='Helper library for consuming the Unmand API',
     url='https://github.com/unmand-systems/unmand-python',
     author='Josiah Khor',

src/unmand/__init__.py

@@ -128,7 +128,7 @@ def queue(self, file_data, guid=None):
 
         return Extraction(None, 'FAILED')
 
-    def poll(self, extraction, max_tries=100, interval=10.0, suppress_output=False):  # pylint: disable=too-many-branches
+    def poll(self, extraction, max_tries=100, interval=10.0, suppress_output=False, with_probabilities=False):  # pylint: disable=too-many-branches
         """Check if prediction is done"""
 
         # Helper function
@@ -144,7 +144,8 @@ def estimate_extraction_time(bounding_box_count):
             time.sleep(interval)
 
             result = requests.get(
-                f'{self.url}extractions/{extraction_guid}/data', auth=TokenAuth(self.token)
+                f'{self.url}extractions/{extraction_guid}/data', auth=TokenAuth(self.token),
+                params={'probabilities': with_probabilities}
             )
 
             if result.status_code == requests.codes.ok:  # pylint: disable=no-member