File size: 3,935 Bytes
3943768 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import wolframalpha
import requests
import os
import argparse
def sanitize_filename(name):
bad_chars = ['[', ']', ',', '/', '\\', '\\w', '\\s', '-', '+', '\"', '\'', '>', '<', ' ', '=', ')', '(', ':', '^']
for char in bad_chars:
name = name.replace(char, "_")
return name
def extract_and_save_images(query, app_id, output_dir):
# Create a client with your app ID
client = wolframalpha.Client(app_id)
# Create output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)
# Send the query
res = client.query(query)
saved_files = []
if res['@success']:
try:
# Print the result
print("<basic_results>")
print(next(res.results).text)
print("</basic_results>")
except StopIteration:
pass
print("\n\n")
print("<detailed_results>")
for i, pod in enumerate(res.pods):
print(f"\nPod: {pod.title}")
for j, sub in enumerate(pod.subpods):
# Print plaintext if available
if sub.plaintext:
print(f" Subpod {j + 1} Text: {sub.plaintext}")
# Save image if available
if hasattr(sub, 'img'):
image_url = sub.img.src
try:
# Download the image
response = requests.get(image_url)
response.raise_for_status()
# Determine the file extension
content_type = response.headers.get('content-type')
ext = content_type.split('/')[-1] if content_type else 'png'
title = sanitize_filename(pod.title)[:20]
sub_title = sanitize_filename(sub.img.title.strip())[:20]
# Create a filename
filename = f"image_{title}_{sub_title}_{i}_{j}.{ext}"
filepath = os.path.join(output_dir, filename)
# Save the image
with open(filepath, 'wb') as f:
f.write(response.content)
saved_files.append(filepath)
print(f" Saved image: {filepath}")
except requests.RequestException as e:
print(f" Error downloading {image_url}: {e}")
print("</detailed_results>")
else:
print(
"Script ran, but query was not successful. Please try a simpler input (e.g. instead of 'plot rule 30', just say 'rule 30') and try again.")
print("Error: ", res['@error'])
return saved_files
def main():
# Set up argument parser
parser = argparse.ArgumentParser(
description="Extract and save images and text from Wolfram Alpha based on a query.")
parser.add_argument("-q", "--query", type=str, required=True, help="The query to send to Wolfram Alpha")
parser.add_argument("-o", "--output_dir", "--file", default="wolfram_images",
help="Output directory for saved images (default: wolfram_images)")
parser.add_argument("-a", "--appid", help="Your Wolfram Alpha App ID")
# Parse arguments
args = parser.parse_args()
# Get App ID from environment variable if not provided as an argument
app_id = args.appid or os.environ.get("WOLFRAM_ALPHA_APPID")
if not app_id:
parser.error(
"Wolfram Alpha App ID must be provided either as an argument or as WOLFRAM_ALPHA_APP_ID environment variable")
try:
print(f"Query: {args.query}\n")
saved_files = extract_and_save_images(args.query, app_id, args.output_dir)
print(f"\nSummary: Saved {len(saved_files)} images to {args.output_dir}/")
except Exception as e:
print(f"An error occurred: {str(e)}")
if __name__ == "__main__":
main()
|