13
13
#
14
14
# You should have received a copy of the GNU Affero General Public License
15
15
# along with this program. If not, see <https://www.gnu.org/licenses/>.
16
- from typing import Match , Union , Any , Type
16
+ from typing import Match , Optional , Tuple , Union , Any , Type
17
17
import operator
18
18
import random
19
19
import math
24
24
from maubot import Plugin , MessageEvent
25
25
from maubot .handlers import command
26
26
27
- pattern_regex = re .compile ("([0-9] {0,9})[dD]([0-9] {1,9})" )
27
+ pattern_regex = re .compile (r"(\d {0,9})d((?:\d|(\{-?\d, *-?\d\})) {1,9})", re . IGNORECASE )
28
28
29
29
_OP_MAP = {
30
30
ast .Add : operator .add ,
@@ -206,37 +206,58 @@ async def roll(self, evt: MessageEvent, pattern: str) -> None:
206
206
207
207
individual_rolls = [] if self .show_rolls else None
208
208
209
- def randomize (number : int , size : int ) -> int :
210
- if size < 0 or number < 0 :
209
+ def randomize (number : int , size : Union [int , Tuple [int , int ]]) -> int :
210
+ size_is_int : bool = isinstance (size , int )
211
+ choices_range : Tuple [int , int ] = (
212
+ size if isinstance (size , tuple ) else (1 , size )
213
+ )
214
+ if not choices_range [0 ] <= choices_range [1 ]:
215
+ raise ValueError (
216
+ "The range's first element must not be greater than the second"
217
+ )
218
+ largest : int = choices_range [1 ]
219
+ if size_is_int and largest < 0 or number < 0 :
211
220
raise ValueError ("randomize() only accepts non-negative values" )
212
- if size == 0 or number == 0 :
221
+ if size_is_int and largest == 0 or number == 0 :
213
222
return 0
214
- elif size == 1 :
223
+ elif size_is_int and largest == 1 :
215
224
return number
216
225
_result = 0
217
226
if number < self .gauss_limit :
218
227
individual = [] if self .show_rolls and number < self .show_rolls_limit else None
219
- for i in range (number ):
220
- roll = random .randint (1 , size )
228
+ for _ in range (number ):
229
+ roll = random .randint (* choices_range )
221
230
if individual is not None :
222
231
individual .append (roll )
223
232
_result += roll
224
- if individual :
225
- individual_rolls .append ((number , size , individual ))
233
+ if individual and individual_rolls is not None :
234
+ individual_rolls .append (
235
+ (
236
+ number ,
237
+ largest
238
+ if size_is_int
239
+ else "{" + ", " .join (str (i ) for i in choices_range ) + "}" ,
240
+ individual
241
+ )
242
+ )
226
243
else :
227
- mean = number * (size + 1 ) / 2
228
- variance = number * (size ** 2 - 1 ) / 12
229
- while _result < number or _result > number * size :
244
+ mean = number * (largest + 1 ) / 2
245
+ variance = number * (largest ** 2 - 1 ) / 12
246
+ while _result < number or _result > number * largest :
230
247
_result = int (random .gauss (mean , math .sqrt (variance )))
231
248
return _result
232
249
233
250
def replacer (match : Match ) -> str :
234
- number = int (match .group (1 ) or "1" )
235
- size = int (match .group (2 ))
251
+ number : int = int (match .group (1 ) or "1" )
252
+ range : Optional [str ] = match .group (3 )
253
+ size : Union [int , Tuple [int , int ]] = (
254
+ int (match .group (2 ))
255
+ if not range
256
+ else tuple (int (i ) for i in range .strip ('{}' ).split (',' ))
257
+ )
236
258
return str (randomize (number , size ))
237
-
238
- pattern = pattern_regex .sub (replacer , pattern )
239
259
try :
260
+ pattern = pattern_regex .sub (replacer , pattern )
240
261
result = Calc .evaluate (pattern )
241
262
if self .round_decimals >= 0 :
242
263
result = round (result , self .round_decimals )
@@ -263,7 +284,8 @@ async def help(self, evt: MessageEvent) -> None:
263
284
"The base command is `!roll`. \\ \n "
264
285
"To roll a dice, pass `XdY` as an argument, where `X` is the "
265
286
"number of dices (optional) and `Y` is the number of sides on "
266
- "each dice. \\ \n "
287
+ "each dice. `Y` can be passed as a specific range as well "
288
+ "(for example: `{0,9}`, `{-5,-1}`). \\ \n "
267
289
"Most Python math and bitwise operators and basic `math` module "
268
290
"functions are also supported, which means you can roll different "
269
291
"kinds of dice and combine the results however you like."
0 commit comments