/
sample.rb
144 lines (124 loc) · 4.37 KB
/
sample.rb
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
143
144
# frozen_string_literal: true
module RuboCop
module Cop
module Style
# This cop is used to identify usages of `shuffle.first`,
# `shuffle.last`, and `shuffle[]` and change them to use
# `sample` instead.
#
# @example
# # bad
# [1, 2, 3].shuffle.first
# [1, 2, 3].shuffle.first(2)
# [1, 2, 3].shuffle.last
# [2, 1, 3].shuffle.at(0)
# [2, 1, 3].shuffle.slice(0)
# [1, 2, 3].shuffle[2]
# [1, 2, 3].shuffle[0, 2] # sample(2) will do the same
# [1, 2, 3].shuffle[0..2] # sample(3) will do the same
# [1, 2, 3].shuffle(random: Random.new).first
#
# # good
# [1, 2, 3].shuffle
# [1, 2, 3].sample
# [1, 2, 3].sample(3)
# [1, 2, 3].shuffle[1, 3] # sample(3) might return a longer Array
# [1, 2, 3].shuffle[1..3] # sample(3) might return a longer Array
# [1, 2, 3].shuffle[foo, bar]
# [1, 2, 3].shuffle(random: Random.new)
class Sample < Cop
MSG = 'Use `%<correct>s` instead of `%<incorrect>s`.'
def_node_matcher :sample_candidate?, <<~PATTERN
(send $(send _ :shuffle $...) ${:first :last :[] :at :slice} $...)
PATTERN
def on_send(node)
sample_candidate?(node) do |shuffle, shuffle_arg, method, method_args|
return unless offensive?(method, method_args)
range = source_range(shuffle, node)
message = message(shuffle_arg, method, method_args, range)
add_offense(node, location: range, message: message)
end
end
def autocorrect(node)
shuffle_node, shuffle_arg, method, method_args =
sample_candidate?(node)
lambda do |corrector|
corrector.replace(source_range(shuffle_node, node),
correction(shuffle_arg, method, method_args))
end
end
private
def offensive?(method, method_args)
case method
when :first, :last
true
when :[], :at, :slice
sample_size(method_args) != :unknown
else
false
end
end
def sample_size(method_args)
case method_args.size
when 1
sample_size_for_one_arg(method_args.first)
when 2
sample_size_for_two_args(*method_args)
end
end
def sample_size_for_one_arg(arg)
if arg.range_type?
range_size(arg)
elsif arg.int_type?
[0, -1].include?(arg.to_a.first) ? nil : :unknown
else
:unknown
end
end
def sample_size_for_two_args(first, second)
return :unknown unless first.int_type? && first.to_a.first.zero?
second.int_type? ? second.to_a.first : :unknown
end
def range_size(range_node) # rubocop:todo Metrics/CyclomaticComplexity
vals = range_node.to_a
return :unknown unless vals.all?(&:int_type?)
low, high = vals.map { |val| val.children[0] }
return :unknown unless low.zero? && high >= 0
case range_node.type
when :erange
(low...high).size
when :irange
(low..high).size
end
end
def source_range(shuffle_node, node)
Parser::Source::Range.new(shuffle_node.source_range.source_buffer,
shuffle_node.loc.selector.begin_pos,
node.source_range.end_pos)
end
def message(shuffle_arg, method, method_args, range)
format(MSG,
correct: correction(shuffle_arg, method, method_args),
incorrect: range.source)
end
def correction(shuffle_arg, method, method_args)
shuffle_arg = extract_source(shuffle_arg)
sample_arg = sample_arg(method, method_args)
args = [sample_arg, shuffle_arg].compact.join(', ')
args.empty? ? 'sample' : "sample(#{args})"
end
def sample_arg(method, method_args)
case method
when :first, :last
extract_source(method_args)
when :[], :slice
sample_size(method_args)
end
end
def extract_source(args)
args.empty? ? nil : args.first.source
end
end
end
end
end