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 # primes [options...] [count...] [numbers...]
  27 #
  28 # Find the first n prime numbers, emitting each one on its own line. When
  29 # given no count, the default is 100 thousand.
  30 #
  31 # All (optional) leading options start with either single or double-dash,
  32 # and most of them change the style/color used. Some of the options are,
  33 # shown in their single-dash form:
  34 #
  35 #     -h          show this help message
  36 #     -help       show this help message
  37 #
  38 #     -check      check all numbers given as cmd-line args for primality
  39 #     -is         check all numbers given as cmd-line args for primality
  40 
  41 
  42 from math import isinf, isnan, modf, sqrt
  43 from sys import argv, exit, stderr, stdout
  44 from typing import List
  45 
  46 
  47 # info is the help message shown when asked to
  48 info = '''
  49 primes [options...] [count...] [numbers...]
  50 
  51 Find the first n prime numbers, emitting each one on its own line. When
  52 given no count, the default is 100 thousand.
  53 
  54 All (optional) leading options start with either single or double-dash,
  55 and most of them change the style/color used. Some of the options are,
  56 shown in their single-dash form:
  57 
  58     -h          show this help message
  59     -help       show this help message
  60 
  61     -check      check all numbers given as cmd-line args for primality
  62     -is         check all numbers given as cmd-line args for primality
  63 '''.strip()
  64 
  65 # handle standard help cmd-line options, quitting right away in that case
  66 if len(argv) > 1 and argv[1] in ('-h', '--h', '-help', '--help'):
  67     print(info, file=stderr)
  68     exit(0)
  69 
  70 
  71 def show_primes(left: int) -> None:
  72     '''Show the number of primes given'''
  73 
  74     # 2 is the only even prime number
  75     if left > 0:
  76         print(2)
  77         left -= 1
  78 
  79     n = 3
  80     while left > 0:
  81         if is_prime(n):
  82             print(n)
  83             left -= 1
  84         n += 2
  85 
  86 
  87 def is_prime(n: int) -> bool:
  88     '''Check if a number is prime, with O(n**0.5) time-complexity'''
  89 
  90     # handle non-sense inputs
  91     if n < 2:
  92         return False
  93 
  94     # 2 is the only even prime number
  95     if n % 2 == 0:
  96         return n == 2
  97 
  98     # check only odd divisors up to the square-root
  99     d = 3
 100     limit = int(sqrt(n))
 101 
 102     while d <= limit:
 103         if n % d == 0:
 104             return False
 105         d += 2
 106     return True
 107 
 108 
 109 def is_float_str(s: str) -> bool:
 110     '''Simplify control-flow for func check_primes'''
 111 
 112     try:
 113         float(s)
 114         return True
 115     except:
 116         return False
 117 
 118 
 119 def check_primes(args: List[str]) -> bool:
 120     '''
 121     Check all the numbers given (as strings) for primality: its output
 122     is lines of tab-separated values (TSV), starting with a header line.
 123     The return value is this func's success, namely its lack of errors.
 124     '''
 125 
 126     errors = 0
 127     print('number\tprime')
 128     for s in args:
 129         try:
 130             f = float(s)
 131             if isnan(f) or isinf(f) or modf(f)[0] != 0.0:
 132                 print(f'{f:,}\tno')
 133             else:
 134                 prime = 'yes' if is_prime(int(f)) else 'no'
 135                 print(f'{f:,}\t{prime}')
 136         except:
 137             print(f'\x1b[31m{s} isn\'t even a number\x1b[0m', file=stderr)
 138             errors += 1
 139     return errors == 0
 140 
 141 
 142 try:
 143     stdout.reconfigure(newline='\n', encoding='utf-8')
 144     check_options = ('check', 'is', 'isprime', 'is-prime', 'is_prime')
 145 
 146     # handle prime-checking mode
 147     if len(argv) > 1 and argv[1].lstrip('-').lower() in check_options:
 148         if len(argv) < 2:
 149             raise ValueError('no numbers given to check')
 150         if not check_primes(argv[2:]):
 151             exit(1)
 152         exit(0)
 153 
 154     # handle main mode, which finds/shows prime numbers: its default count
 155     # is 100 thousand, unless a leading cmd-line arg overrides it
 156     n = int(1e5)
 157     if len(argv) > 1:
 158         try:
 159             n = int(float(argv[1]))
 160         except:
 161             raise ValueError(f'unsupported leading option {argv[1]}')
 162 
 163     print(f'showing the first {n:,} primes', file=stderr)
 164     # running inside a func speeds things up in older versions of python
 165     show_primes(n)
 166 except KeyboardInterrupt:
 167     print('script was force-quit early', file=stderr)
 168     exit(1)
 169 except Exception as e:
 170     print(f'\x1b[31m{e}\x1b[0m', file=stderr)
 171     exit(1)