File: geo_commands.rb

package info (click to toggle)
ruby-fakeredis 0.8.0-7
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 628 kB
  • sloc: ruby: 4,868; makefile: 2
file content (142 lines) | stat: -rw-r--r-- 3,980 bytes parent folder | download | duplicates (2)
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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
require "fakeredis/geo_set"

module FakeRedis
  module GeoCommands
    DISTANCE_UNITS = {
      "m" => 1,
      "km" => 1000,
      "ft" => 0.3048,
      "mi" => 1609.34
    }

    REDIS_DOUBLE_PRECISION = 4
    REDIS_GEOHASH_SIZE = 10

    def geoadd(key, *members)
      raise_argument_error("geoadd") if members.empty? || members.size % 3 != 0

      set = (data[key] ||= GeoSet.new)
      prev_size = set.size
      members.each_slice(3) do |member|
        set.add(*member)
      end
      set.size - prev_size
    end

    def geodist(key, member1, member2, unit = "m")
      unit = unit.to_s
      raise_command_error("ERR unsupported unit provided. please use #{DISTANCE_UNITS.keys.join(', ')}") unless DISTANCE_UNITS.include?(unit)

      set = (data[key] || GeoSet.new)
      point1 = set.get(member1)
      point2 = set.get(member2)
      if point1 && point2
        distance = point1.distance_to(point2)
        distance_in_units = distance / DISTANCE_UNITS[unit]
        distance_in_units.round(REDIS_DOUBLE_PRECISION).to_s
      end
    end

    def geohash(key, member)
      members = Array(member)
      raise_argument_error("geohash") if members.empty?
      set = (data[key] || GeoSet.new)
      members.map do |member|
        point = set.get(member)
        point.geohash(REDIS_GEOHASH_SIZE) if point
      end
    end

    def geopos(key, member)
      return nil unless data[key]

      members = Array(member)
      set = (data[key] || GeoSet.new)
      members.map do |member|
        point = set.get(member)
        [point.lon.to_s, point.lat.to_s] if point
      end
    end

    def georadius(*args)
      args = args.dup
      raise_argument_error("georadius") if args.size < 5
      key, lon, lat, radius, unit, *rest = args
      raise_argument_error("georadius") unless DISTANCE_UNITS.has_key?(unit)
      radius *= DISTANCE_UNITS[unit]

      set = (data[key] || GeoSet.new)
      center = GeoSet::Point.new(lon, lat, nil)

      do_georadius(set, center, radius, unit, rest)
    end

    def georadiusbymember(*args)
      args = args.dup
      raise_argument_error("georadiusbymember") if args.size < 4
      key, member, radius, unit, *rest = args
      raise_argument_error("georadiusbymember") unless DISTANCE_UNITS.has_key?(unit)
      radius *= DISTANCE_UNITS[unit]

      set = (data[key] || GeoSet.new)
      center = set.get(member)
      raise_command_error("ERR could not decode requested zset member") unless center

      do_georadius(set, center, radius, unit, args)
    end

    private

    def do_georadius(set, center, radius, unit, args)
      points = set.points_within_radius(center, radius)

      options = georadius_options(args)

      if options[:asc]
        points.sort_by! { |p| p.distance_to(center) }
      elsif options[:desc]
        points.sort_by! { |p| -p.distance_to(center) }
      end

      points = points.take(options[:count]) if options[:count]
      extras = options[:extras]
      return points.map(&:name) if extras.empty?

      points.map do |point|
        member = [point.name]

        extras.each do |extra|
          case extra
          when "WITHCOORD"
            member << [point.lon.to_s, point.lat.to_s]
          when "WITHDIST"
            distance = point.distance_to(center)
            distance_in_units = distance / DISTANCE_UNITS[unit]
            member << distance_in_units.round(REDIS_DOUBLE_PRECISION).to_s
          when "WITHHASH"
            member << point.geohash(REDIS_GEOHASH_SIZE)
          end
        end

        member
      end
    end

    def georadius_options(args)
      options = {}
      args = args.map { |arg| arg.to_s.upcase }

      if idx = args.index("COUNT")
        options[:count] = Integer(args[idx + 1])
      end

      options[:asc]  = true if args.include?("ASC")
      options[:desc] = true if args.include?("DESC")

      extras = args & ["WITHCOORD", "WITHDIST", "WITHHASH"]
      options[:extras] = extras

      options
    end
  end
end