File: primes.py
   1 #!/usr/bin/python3
   2 
   3 # The MIT License (MIT)
   4 #
   5 # Copyright © 2024 pacman64
   6 #
   7 # Permission is hereby granted, free of charge, to any person obtaining a copy
   8 # of this software and associated documentation files (the “Software”), to deal
   9 # in the Software without restriction, including without limitation the rights
  10 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  11 # copies of the Software, and to permit persons to whom the Software is
  12 # furnished to do so, subject to the following conditions:
  13 #
  14 # The above copyright notice and this permission notice shall be included in
  15 # all copies or substantial portions of the Software.
  16 #
  17 # THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  18 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  19 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  20 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  21 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  22 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  23 # SOFTWARE.
  24 
  25 
  26 from math import isinf, isnan, modf, sqrt
  27 from sys import argv, exit, stderr
  28 from typing import List
  29 
  30 
  31 info = '''
  32 primes [options...] [count...] [numbers...]
  33 
  34 
  35 Find or check prime numbers.
  36 
  37 By default it finds the first n prime numbers, emitting each one on its own
  38 line, with a default of 100 thousand when not given a count.
  39 
  40 When using one of the check-prime options, all later cmd-line arguments are
  41 expected to be numbers, and are checked for primality.
  42 
  43 All (optional) leading options start with either single or double-dash, or
  44 even no dashes at all. Some of the options are, shown in their single-dash
  45 form:
  46 
  47     -h          show this help message
  48     -help       show this help message
  49 
  50     -check      check all numbers given as cmd-line args for primality
  51     -isprime    same as option -check
  52     -is-prime   same as option -check
  53     -is         same as option -check
  54 '''
  55 
  56 # handle standard help cmd-line options, quitting right away in that case
  57 if len(argv) > 1 and argv[1].lstrip('-') in ('h', 'help'):
  58     print(info.strip(), file=stderr)
  59     exit(0)
  60 
  61 
  62 def show_primes(left: int) -> None:
  63     'Show the number of primes given.'
  64 
  65     # 2 is the only even prime number
  66     if left > 0:
  67         print(2)
  68         left -= 1
  69 
  70     n = 3
  71     while left > 0:
  72         if is_prime(n):
  73             print(n)
  74             left -= 1
  75         n += 2
  76 
  77 
  78 def is_prime(n: int) -> bool:
  79     'Check if a number is prime, with O(n**0.5) time-complexity.'
  80 
  81     # handle non-sense inputs
  82     if n < 2:
  83         return False
  84 
  85     # 2 is the only even prime number
  86     if n % 2 == 0:
  87         return n == 2
  88 
  89     # check only odd divisors up to the square-root
  90     d = 3
  91     limit = int(sqrt(n))
  92 
  93     while d <= limit:
  94         if n % d == 0:
  95             return False
  96         d += 2
  97     return True
  98 
  99 
 100 def is_float_str(s: str) -> bool:
 101     'Simplify control-flow for func check_primes.'
 102 
 103     try:
 104         float(s)
 105         return True
 106     except Exception:
 107         return False
 108 
 109 
 110 def check_primes(args: List[str]) -> bool:
 111     '''
 112     Check all the numbers given (as strings) for primality: its output
 113     is lines of tab-separated values (TSV), starting with a header line.
 114     The return value is this func's success, namely its lack of errors.
 115     '''
 116 
 117     errors = 0
 118     print('number\tprime')
 119     for s in args:
 120         try:
 121             f = float(s)
 122             if isnan(f) or isinf(f) or modf(f)[0] != 0.0:
 123                 print(f'{f:,}\tno')
 124             else:
 125                 prime = 'yes' if is_prime(int(f)) else 'no'
 126                 print(f'{f:,}\t{prime}')
 127         except Exception:
 128             print(f'\x1b[31m{s} isn\'t even a number\x1b[0m', file=stderr)
 129             errors += 1
 130     return errors == 0
 131 
 132 
 133 try:
 134     check_options = ('check', 'is', 'isprime', 'is-prime', 'is_prime')
 135 
 136     # handle prime-checking mode
 137     if len(argv) > 1 and argv[1].lstrip('-').lower() in check_options:
 138         if len(argv) < 2:
 139             raise ValueError('no numbers given to check')
 140         if not check_primes(argv[2:]):
 141             exit(1)
 142         exit(0)
 143 
 144     # handle main mode, which finds/shows prime numbers: its default count
 145     # is 100 thousand, unless a leading cmd-line arg overrides it
 146     n = int(1e5)
 147     if len(argv) > 1:
 148         try:
 149             n = int(float(argv[1]))
 150         except Exception:
 151             raise ValueError(f'unsupported leading option {argv[1]}')
 152 
 153     print(f'showing the first {n:,} primes', file=stderr)
 154     # running inside a func speeds things up in older versions of python
 155     show_primes(n)
 156 except BrokenPipeError:
 157     # quit quietly, instead of showing a confusing error message
 158     stderr.close()
 159 except KeyboardInterrupt:
 160     print('\x1b[31mscript was force-quit early\x1b[0m', file=stderr)
 161     exit(2)
 162 except Exception as e:
 163     print(f'\x1b[31m{e}\x1b[0m', file=stderr)
 164     exit(1)