datapungibea 2.0.0__py3-none-any.whl

datapungibea/codegen.py

@@ -0,0 +1,93 @@
+r'''
+   Code snippet generation for reproducing queries with requests directly.
+   Simplified to always use environment variable pattern.
+'''
+
+import json
+import logging
+from textwrap import dedent
+from typing import Any, Dict, Optional
+
+import pyperclip
+
+logger = logging.getLogger(__name__)
+
+
+class transformIncludeCodeSnippet:
+    def transformIncludeCodeSnippet(
+        self,
+        query: Dict[str, Any],
+        baseRequest: Dict[str, Any],
+        pandasCode: str = "",
+    ) -> str:
+        '''
+        Generate a reproducible code snippet for a query.
+        
+        Always generates code that uses environment variable pattern:
+            import os
+            key = os.getenv("API_KEY_BEA")
+        
+        Args:
+            query: the query dict with url and params
+            baseRequest: the base request dict
+            pandasCode: additional pandas code for data transformation
+            
+        Returns:
+            Complete Python code string that reproduces the query
+        '''
+        apiCode = self.getApiCode(query)
+        queryCode = self.getQueryCode(query, baseRequest, pandasCode)
+        return apiCode + queryCode
+
+    def getQueryCode(self, query: Dict[str, Any], baseRequest: Dict[str, Any], pandasCode: str = "") -> str:
+        '''
+        Generate the query execution code.
+        '''
+        queryClean = {'url': query['url'], 'params': query['params']}
+        queryClean['url'] = 'url'
+
+        # Replace API key with variable
+        if 'UserID' in queryClean['params'] and 'UserID' in baseRequest.get('params', {}):
+            queryClean['params']['UserID'] = '{}'
+
+        queryCode = dedent('''\
+            query = {}
+            retrivedData = requests.get(**query)
+
+            {} # replace json by xml if this is the request format
+        ''').format(json.dumps(queryClean), pandasCode)
+
+        queryCode = queryCode.replace('"url": "url"', '"url": url')
+        queryCode = queryCode.replace('"UserID": "{}"', '"UserID": key')
+
+        return queryCode
+
+    def getApiCode(self, query: Dict[str, Any]) -> str:
+        '''
+        Generate the API setup code using environment variable pattern.
+        '''
+        url = query.get('url', 'https://apps.bea.gov/api/data/')
+        
+        code = dedent('''\
+            import requests
+            import os
+            import pandas as pd
+
+            url = "{url}"
+            key = os.getenv("API_KEY_BEA")
+        ''').format(url=url)
+
+        return code
+
+    def clipcode(self, lastLoad: Dict[str, Any]) -> None:
+        r'''
+           Copy the string to the user's clipboard
+        '''
+        try:
+            pyperclip.copy(lastLoad['code'])
+            logger.info("Code copied to clipboard")
+        except (KeyError, TypeError):
+            logger.warning(
+                "Loaded session does not have a code entry. "
+                "Re-run with verbose option set to True, e.g.: data.NIPA(..., verbose=True)"
+            )