Skip to main content

redis/commands/
geo.rs

1//! Defines types to use with the geospatial commands.
2
3use crate::errors::{ParsingError, invalid_type_error};
4use crate::types::{FromRedisValue, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value};
5
6/// Units used by [`geo_dist`][1] and [`geo_radius`][2].
7///
8/// [1]: ../trait.Commands.html#method.geo_dist
9/// [2]: ../trait.Commands.html#method.geo_radius
10#[non_exhaustive]
11pub enum Unit {
12    /// Represents meters.
13    Meters,
14    /// Represents kilometers.
15    Kilometers,
16    /// Represents miles.
17    Miles,
18    /// Represents feet.
19    Feet,
20}
21
22impl ToRedisArgs for Unit {
23    fn write_redis_args<W>(&self, out: &mut W)
24    where
25        W: ?Sized + RedisWrite,
26    {
27        let unit = match *self {
28            Self::Meters => "m",
29            Self::Kilometers => "km",
30            Self::Miles => "mi",
31            Self::Feet => "ft",
32        };
33        out.write_arg(unit.as_bytes());
34    }
35}
36
37impl ToSingleRedisArg for Unit {}
38
39/// A coordinate (longitude, latitude).
40///
41/// Can be used with [`geo_pos`][1] to parse response from Redis.
42///
43/// [1]: ../trait.Commands.html#method.geo_pos
44///
45/// `T` is the type of the every value.
46///
47/// * You may want to use either `f64` or `f32` if you want to perform mathematical operations.
48/// * To keep the raw value from Redis, use `String`.
49#[allow(clippy::derive_partial_eq_without_eq)] // allow f32/f64 here, which don't implement Eq
50#[derive(Debug, PartialEq)]
51pub struct Coord<T> {
52    /// Longitude
53    pub longitude: T,
54    /// Latitude
55    pub latitude: T,
56}
57
58impl<T> Coord<T> {
59    /// Create a new Coord with the (longitude, latitude)
60    pub fn lon_lat(longitude: T, latitude: T) -> Self {
61        Self {
62            longitude,
63            latitude,
64        }
65    }
66}
67
68impl<T: FromRedisValue> FromRedisValue for Coord<T> {
69    fn from_redis_value_ref(v: &Value) -> Result<Self, ParsingError> {
70        let values: Vec<T> = FromRedisValue::from_redis_value_ref(v)?;
71        let mut values = values.into_iter();
72        let (longitude, latitude) = match (values.next(), values.next(), values.next()) {
73            (Some(longitude), Some(latitude), None) => (longitude, latitude),
74            _ => invalid_type_error!(v, "Expect a pair of numbers"),
75        };
76        Ok(Self {
77            longitude,
78            latitude,
79        })
80    }
81
82    fn from_redis_value(v: Value) -> Result<Self, ParsingError> {
83        Self::from_redis_value_ref(&v)
84    }
85}
86
87impl<T: ToRedisArgs> ToRedisArgs for Coord<T> {
88    fn write_redis_args<W>(&self, out: &mut W)
89    where
90        W: ?Sized + RedisWrite,
91    {
92        ToRedisArgs::write_redis_args(&self.longitude, out);
93        ToRedisArgs::write_redis_args(&self.latitude, out);
94    }
95
96    fn num_of_args(&self) -> usize {
97        2
98    }
99}
100
101/// Options to sort results from [GEORADIUS][1] and [GEORADIUSBYMEMBER][2] commands
102///
103/// [1]: https://redis.io/commands/georadius
104/// [2]: https://redis.io/commands/georadiusbymember
105#[derive(Default)]
106#[non_exhaustive]
107pub enum RadiusOrder {
108    /// Don't sort the results
109    #[default]
110    Unsorted,
111
112    /// Sort returned items from the nearest to the farthest, relative to the center.
113    Asc,
114
115    /// Sort returned items from the farthest to the nearest, relative to the center.
116    Desc,
117}
118
119/// Options for the [GEORADIUS][1] and [GEORADIUSBYMEMBER][2] commands
120///
121/// [1]: https://redis.io/commands/georadius
122/// [2]: https://redis.io/commands/georadiusbymember
123///
124/// # Example
125///
126/// ```rust,no_run
127/// use redis::{Commands, RedisResult};
128/// use redis::geo::{RadiusSearchResult, RadiusOptions, RadiusOrder, Unit};
129/// fn nearest_in_radius(
130///     con: &mut redis::Connection,
131///     key: &str,
132///     longitude: f64,
133///     latitude: f64,
134///     meters: f64,
135///     limit: usize,
136/// ) -> RedisResult<Vec<RadiusSearchResult>> {
137///     let opts = RadiusOptions::default()
138///         .order(RadiusOrder::Asc)
139///         .limit(limit);
140///     con.geo_radius(key, longitude, latitude, meters, Unit::Meters, opts)
141/// }
142/// ```
143#[derive(Default)]
144pub struct RadiusOptions {
145    with_coord: bool,
146    with_dist: bool,
147    count: Option<usize>,
148    order: RadiusOrder,
149    store: Option<Vec<Vec<u8>>>,
150    store_dist: Option<Vec<Vec<u8>>>,
151}
152
153impl RadiusOptions {
154    /// Limit the results to the first N matching items.
155    pub fn limit(mut self, n: usize) -> Self {
156        self.count = Some(n);
157        self
158    }
159
160    /// Return the distance of the returned items from the specified center.
161    /// The distance is returned in the same unit as the unit specified as the
162    /// radius argument of the command.
163    pub fn with_dist(mut self) -> Self {
164        self.with_dist = true;
165        self
166    }
167
168    /// Return the `longitude, latitude` coordinates of the matching items.
169    pub fn with_coord(mut self) -> Self {
170        self.with_coord = true;
171        self
172    }
173
174    /// Sort the returned items
175    pub fn order(mut self, o: RadiusOrder) -> Self {
176        self.order = o;
177        self
178    }
179
180    /// Store the results in a sorted set at `key`, instead of returning them.
181    ///
182    /// This feature can't be used with any `with_*` method.
183    pub fn store<K: ToRedisArgs>(mut self, key: K) -> Self {
184        self.store = Some(ToRedisArgs::to_redis_args(&key));
185        self
186    }
187
188    /// Store the results in a sorted set at `key`, with the distance from the
189    /// center as its score. This feature can't be used with any `with_*` method.
190    pub fn store_dist<K: ToRedisArgs>(mut self, key: K) -> Self {
191        self.store_dist = Some(ToRedisArgs::to_redis_args(&key));
192        self
193    }
194}
195
196impl ToRedisArgs for RadiusOptions {
197    fn write_redis_args<W>(&self, out: &mut W)
198    where
199        W: ?Sized + RedisWrite,
200    {
201        if self.with_coord {
202            out.write_arg(b"WITHCOORD");
203        }
204
205        if self.with_dist {
206            out.write_arg(b"WITHDIST");
207        }
208
209        if let Some(n) = self.count {
210            out.write_arg(b"COUNT");
211            out.write_arg_fmt(n);
212        }
213
214        match self.order {
215            RadiusOrder::Asc => out.write_arg(b"ASC"),
216            RadiusOrder::Desc => out.write_arg(b"DESC"),
217            _ => (),
218        }
219
220        if let Some(ref store) = self.store {
221            out.write_arg(b"STORE");
222            for i in store {
223                out.write_arg(i);
224            }
225        }
226
227        if let Some(ref store_dist) = self.store_dist {
228            out.write_arg(b"STOREDIST");
229            for i in store_dist {
230                out.write_arg(i);
231            }
232        }
233    }
234
235    fn num_of_args(&self) -> usize {
236        let mut n: usize = 0;
237        if self.with_coord {
238            n += 1;
239        }
240        if self.with_dist {
241            n += 1;
242        }
243        if self.count.is_some() {
244            n += 2;
245        }
246        match self.order {
247            RadiusOrder::Asc | RadiusOrder::Desc => n += 1,
248            _ => {}
249        }
250        n += 1 + self.store.as_ref().map(|v| v.len()).unwrap_or(0);
251        n += 1 + self.store_dist.as_ref().map(|v| v.len()).unwrap_or(0);
252        n
253    }
254}
255
256/// Contain an item returned by [`geo_radius`][1] and [`geo_radius_by_member`][2].
257///
258/// [1]: ../trait.Commands.html#method.geo_radius
259/// [2]: ../trait.Commands.html#method.geo_radius_by_member
260pub struct RadiusSearchResult {
261    /// The name that was found.
262    pub name: String,
263    /// The coordinate if available.
264    pub coord: Option<Coord<f64>>,
265    /// The distance if available.
266    pub dist: Option<f64>,
267}
268
269impl FromRedisValue for RadiusSearchResult {
270    fn from_redis_value(v: Value) -> Result<Self, ParsingError> {
271        match v {
272            Value::BulkString(b) => {
273                let s = String::from_utf8(b)?;
274                Ok(Self {
275                    name: s,
276                    coord: None,
277                    dist: None,
278                })
279            }
280            Value::Array(items) => Self::parse_multi_values(items),
281            _ => invalid_type_error!(v, "Response type not RadiusSearchResult compatible."),
282        }
283    }
284}
285
286impl RadiusSearchResult {
287    fn parse_multi_values(items: Vec<Value>) -> Result<Self, ParsingError> {
288        let mut iter = items.into_iter();
289
290        // First item is always the member name
291        let name: String = match iter.next().map(FromRedisValue::from_redis_value) {
292            Some(Ok(n)) => n,
293            _ => return Err(arcstr::literal!("Missing member name").into()),
294        };
295
296        let (dist, coord) = match (iter.next(), iter.next()) {
297            (None, None) => (None, None),
298            (Some(Value::Array(coords)), None) => {
299                (None, Some(Coord::from_redis_value(Value::Array(coords))?))
300            }
301            (Some(dist), coord) => {
302                let dist = FromRedisValue::from_redis_value(dist)?;
303
304                let coord = match coord.map(FromRedisValue::from_redis_value) {
305                    Some(Ok(c)) => Some(c),
306                    _ => None,
307                };
308
309                (dist, coord)
310            }
311            _ => invalid_type_error!("Response type not RadiusSearchResult compatible."),
312        };
313
314        Ok(Self { name, coord, dist })
315    }
316}
317
318#[cfg(test)]
319mod tests {
320    use super::{Coord, RadiusOptions, RadiusOrder};
321    use crate::types::ToRedisArgs;
322    use std::str;
323
324    macro_rules! assert_args {
325        ($value:expr, $($args:expr),+) => {
326            let args = $value.to_redis_args();
327            let strings: Vec<_> = args.iter()
328                                      .map(|a| str::from_utf8(a.as_ref()).unwrap())
329                                      .collect();
330            assert_eq!(strings, vec![$($args),+]);
331        }
332    }
333
334    #[test]
335    fn test_coord_to_args() {
336        let member = ("Palermo", Coord::lon_lat("13.361389", "38.115556"));
337        assert_args!(&member, "Palermo", "13.361389", "38.115556");
338    }
339
340    #[test]
341    fn test_radius_options() {
342        // Without options, should not generate any argument
343        let empty = RadiusOptions::default();
344        assert_eq!(ToRedisArgs::to_redis_args(&empty).len(), 0);
345
346        // Some combinations with WITH* options
347        let opts = RadiusOptions::default;
348
349        assert_args!(opts().with_coord().with_dist(), "WITHCOORD", "WITHDIST");
350
351        assert_args!(opts().limit(50), "COUNT", "50");
352
353        assert_args!(opts().limit(50).store("x"), "COUNT", "50", "STORE", "x");
354
355        assert_args!(
356            opts().limit(100).store_dist("y"),
357            "COUNT",
358            "100",
359            "STOREDIST",
360            "y"
361        );
362
363        assert_args!(
364            opts().order(RadiusOrder::Asc).limit(10).with_dist(),
365            "WITHDIST",
366            "COUNT",
367            "10",
368            "ASC"
369        );
370    }
371}