~handlerug/handlebot

ref: 7a899986d3a7d250418ed425c03a9e5391820979 handlebot/wolframalpha/wolframalpha.go -rw-r--r-- 2.8 KiB
7a899986Umar Getagazov Set user agent 3 months ago
                                                                                
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
package wolframalpha

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"
)

type Result struct {
	Success     bool        `json:"success"`
	Input       string      `json:"inputstring"`
	Pods        []Pod       `json:"pods"`
	DidYouMeans DidYouMeans `json:"didyoumeans"`
}

type DidYouMean struct {
	Score float32 `json:"score,string"`
	Value string  `json:"val"`
}

type DidYouMeans []DidYouMean

func (d DidYouMeans) Len() int           { return len(d) }
func (d DidYouMeans) Swap(i, j int)      { d[i], d[j] = d[j], d[i] }
func (d DidYouMeans) Less(i, j int) bool { return d[i].Score > d[j].Score }

func (d DidYouMeans) String() string {
	var b strings.Builder
	for i, val := range d {
		b.WriteString(val.Value)
		if i != len(d)-1 {
			b.WriteString(", ")
		}
	}
	return b.String()
}

func (d *DidYouMeans) UnmarshalJSON(b []byte) error {
	switch b[0] {
	case '{':
		return d.unmarshalSingle(b)
	case '[':
		return d.unmarshalMany(b)
	}
	return fmt.Errorf("invalid JSON data; did-you-means must be either an " +
		"object or an array of objects")
}

func (d *DidYouMeans) unmarshalSingle(b []byte) error {
	var single DidYouMean
	if err := json.Unmarshal(b, &single); err != nil {
		return err
	}
	*d = DidYouMeans{single}
	return nil
}

func (d *DidYouMeans) unmarshalMany(b []byte) error {
	many := []DidYouMean{}
	if err := json.Unmarshal(b, &many); err != nil {
		return err
	}
	*d = many
	return nil
}

type Pod struct {
	ID      string   `json:"id"`
	Subpods []Subpod `json:"subpods"`
}

type Subpod struct {
	PlainText string `json:"plaintext"`
}

type Client struct {
	AppID      string
	HTTPClient *http.Client
}

func (c *Client) Query(ctx context.Context, input string) (*Result, error) {
	q := url.Values{}
	q.Set("input", input)
	q.Set("appid", c.AppID)
	q.Set("output", "JSON")
	q.Add("includepodid", "Input")
	q.Add("includepodid", "Result")
	// I would have included reinterpret=true as well, but it doesn't seem to
	// affect responses at all. Besides, it also removes didyoumeans even in
	// the case it fails, which are quite useful.

	req, _ := http.NewRequestWithContext(ctx, "GET", "http://api.wolframalpha.com/v2/query", nil)
	req.URL.RawQuery = q.Encode()

	resp, err := c.httpClient().Do(req)
	if err != nil {
		if err.(*url.Error).Timeout() {
			return nil, &TimeoutError{}
		} else {
			return nil, &RequestError{err}
		}
	}
	defer resp.Body.Close()

	var apiResp struct {
		Result Result `json:"queryresult"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&apiResp); err != nil {
		return nil, &ResponseError{err}
	}

	if !apiResp.Result.Success {
		return nil, &NoResultsError{
			DidYouMeans: apiResp.Result.DidYouMeans,
		}
	}
	return &apiResp.Result, nil
}

func (c *Client) httpClient() *http.Client {
	if c.HTTPClient != nil {
		return c.HTTPClient
	}
	return http.DefaultClient
}