Fix handling of AI extract rules

Author: daolf

README.md

@@ -40,6 +40,8 @@ Signup to ScrapingBee to [get your API key](https://app.scrapingbee.com/account/
         'device': 'desktop',
         # Use some data extraction rules
         'extract_rules': {'title': 'h1'},
+        # Use AI to extract data from the page
+        'ai_extract_rules': {'product_name': 'The name of the product', 'price': 'The price in USD'},
         # Wrap response in JSON
         'json_response': False,
         # Interact with the webpage you want to scrape 

scrapingbee/__version__.py

@@ -1 +1 @@
-__version__ = "2.0.1"
+__version__ = "2.0.2"

scrapingbee/utils.py

@@ -46,6 +46,8 @@ def process_params(params: dict) -> dict:
             new_params[k] = process_cookies(v)
         elif k == 'extract_rules':
             new_params[k] = process_json_stringify_param(v, 'extract_rules')
+        elif k == 'ai_extract_rules':
+            new_params[k] = process_json_stringify_param(v, 'ai_extract_rules')
         elif k == 'js_scenario':
             new_params[k] = process_json_stringify_param(v, 'js_scenario')
         else:

tests/test_client.py

@@ -113,6 +113,27 @@ def test_get_with_js_scenario(mock_session, client):
     )
 
 
+@mock.patch('scrapingbee.client.Session')
+def test_get_with_ai_extract_rules(mock_session, client):
+    '''It should format the ai_extract_rules and add them to the url'''
+    client.get('https://httpbin.org', params={
+        'ai_extract_rules': {
+            "product_name": "The name of the product",
+            "price": "The price in USD"
+        }
+    })
+
+    mock_session.return_value.request.assert_called_with(
+        'GET',
+        'https://app.scrapingbee.com/api/v1/'
+        '?api_key=API_KEY&url=https%3A%2F%2Fhttpbin.org&'
+        'ai_extract_rules=%7B%22product_name%22%3A+%22The+name+of+the+product%22%2C+%22'
+        'price%22%3A+%22The+price+in+USD%22%7D',
+        data=None,
+        headers=DEFAULT_HEADERS,
+    )
+
+
 @mock.patch('scrapingbee.client.Session')
 def test_post(mock_session, client):
     '''It should make a POST request with some data'''

tests/test_utils.py

@@ -46,6 +46,12 @@ def test_process_js_scenario():
     assert output == '{"instructions": [{"click": "#buttonId"}]}'
 
 
+def test_process_ai_extract_rules():
+    """It should format ai_extract_rules to a stringified JSON"""
+    output = process_json_stringify_param({"product_name": "The name of the product", "price": "The price in USD"}, "ai_extract_rules")
+    assert output == '{"product_name": "The name of the product", "price": "The price in USD"}'
+
+
 def test_process_params():
     """It should keep boolean parameters"""
     output = process_params({"render_js": True})