A python script for bookkeeping with notion in terminal!

bkp.py

#!/usr/bin/env python3
"""This is a python script for bookkeeping with notion in CLI.

Author:   AlanDecode
Link:     https://blog.imalan.cn/archives/track-your-expense-with-notion/
License:  MIT
Version:  1.0

First, you need to setup a notion database. Create a database(Table-Full page),
and set its properties(columns) as:

- 描述, Title
- 金额, Number, you can format it as Yuan or something
- 交易标签, Multi-select
- 交易日期, Date, no need to include time

Open a notion webpage which has been logged in, press F12, in
Application - Cookies - https://www.notion.so, find token_v2.

Add following contents to ~/.bkprc:

```
BKP_TOKEN_V2 = the token_v2 from notion cookies
BKP_DATABASE_URL = the database url
```

Or set them in environment variables:

```
export BKP_TOKEN_V2=...
export BKP_DATABASE_URL=...
```

We need notion-py to access notion database, so try `pip3 install notion` first,
then just run `python3 bkp.py` to use this script.

Or you can install it to somewhere in PATH:

```
chmod u+x bkp.py
sudo ln -s $(pwd)/bkp.py /usr/local/bin/bkp
```

Then you can call `bkp` in terminal to use this script anywhere.

Input hints while using this script:

- Description: The description of transaction, string
-      Amount: The amount of money of this transaction, float. Or you can input
               something like 32*2+15/2, try it!
-        Tags: The tags of this transaction, string separated by , or ，or 、
-        Date: The transaction date, like `2021-05-12`. For convenience, press
               enter for today, -1 for yesterday, etc.
"""
import datetime
import logging
import os

logging.basicConfig(level=logging.INFO, format='=> %(message)s')
logger = logging.getLogger('BKP')


class Color:
  BLACK = 30
  RED = 31
  GREEN = 32
  YELLOW = 33
  BLUE = 34
  MAGENTA = 35
  CYAN = 36
  WHITE = 37


def clr(text: str, fg: int = Color.GREEN) -> str:
  """Returns colored text.
  """
  return f'\033[{fg}m{text}\033[0m'


def red(text: str) -> str:
  return clr(text, Color.RED)


def green(text: str) -> str:
  return clr(text, Color.GREEN)


def blue(text: str) -> str:
  return clr(text, Color.BLUE)


def yellow(text: str) -> str:
  return clr(text, Color.YELLOW)


def get_config() -> dict:
  """Get configs from file and env.
  """
  config = {}

  # Load config from ~/.bkprc file, if exist
  bkprc_path = os.path.expanduser('~/.bkprc')
  logger.info(clr(f'Loading config from {bkprc_path}', Color.BLUE))
  if os.path.exists(bkprc_path) and os.path.isfile(bkprc_path):
    with open(bkprc_path, 'r') as f:
      for line in f.readlines():
        if not line.strip():
          continue
        fields = line.split('=')
        k, v = fields[0], '='.join(fields[1:])
        config[k.strip()] = v.strip()

  # Load config from env
  logger.info(clr(f'Loading config from ENV', Color.BLUE))
  for env_name, env_value in os.environ.items():
    if env_name.startswith('BKP_'):
      config[env_name] = env_value

  if 'BKP_TOKEN_V2' not in config or 'BKP_DATABASE_URL' not in config:
    logger.error(red(
        f'Please set `BKP_TOKEN_V2` and `BKP_DATABASE_URL` in '
        f'{os.path.expanduser("~/.bkprc")} or environment variables.'))
    exit(-1)

  return config


try:
  from notion.client import NotionClient
except ImportError:
  logger.error(
      clr('We need notion-py! Try `pip install notion` first!', Color.RED))
  exit(-1)


def submit_to_notion(cv: NotionClient, desc: str, amount: float, tags: list,
                     date: datetime.date):
  """Create an new record and submit to notion.
  """
  row = cv.collection.add_row()
  row.miao_shu = desc
  row.jin_e = amount
  row.jiao_yi_biao_qian = tags
  row.jiao_yi_ri_qi = date


def get_amount(prompt: str):
  """Get eval results from input.
  """
  try:
    return float(eval(input(prompt)))
  except Exception:
    return None


def main():
  """Main loop.
  """
  cfg = get_config()

  # Init notion instance
  try:
    logger.info(blue('Connecting to Notion...'))
    client = NotionClient(token_v2=cfg['BKP_TOKEN_V2'])
    cv = client.get_collection_view(cfg['BKP_DATABASE_URL'])
  except KeyError as e:
    logger.error(red(f'Please set {e}.'))
    exit(-1)
  except Exception as e:
    logger.error(red('Can not connect to Notion, is your token expired?'))
    exit(-1)

  # Output current status
  aggregations = [{
      "property": "jin_e",
      "aggregator": "count",
      "id": "num_entries",
  }]
  num_entries = cv.build_query(
      aggregations=aggregations).execute().get_aggregate("num_entries")
  logger.info(
      blue(
          f'Awesome! You already have {num_entries} records, let\'s add more!'))

  while True:
    # Main loop, always wait for input
    try:
      print(yellow('----------------'))
      desc = input(yellow('|-> Description: '))

      # Parse transaction amount
      amount = None
      while True:
        amount = get_amount(yellow('|->      Amount: '))
        if amount is not None:
          break
        else:
          logger.error(red('Can\'t parse amount, please enter again!'))

      # Parse tag list
      tags = input(yellow('|->        Tags: '))
      for sep in ('，', '、'):
        tags = tags.replace(sep, ',')
      tags = set([item.strip() for item in tags.split(',') if item.strip()])

      # Parse transaction date
      date = input(yellow('|->        Date: '))
      print(yellow('----------------'))
      if date == '':
        date = 0
      try:
        # Parse date as relative date from today
        date = datetime.date.today() + datetime.timedelta(days=int(date))
      except ValueError:
        # Parse date in year-month-day format
        date = date.strip().split(' ')[0]
        for sep in ('：', ':', '/', '.', '—'):
          date = date.replace(sep, '-')
        date = '-'.join([item.strip() for item in date.split('-')])
        date = datetime.datetime.strptime(date, '%Y-%m-%d')
        date = datetime.date(date.year, date.month, date.day)
    except Exception as e:
      logger.error(red('Be careful with what you entered!'))
      print(str(e))
      continue

    # Add entry to notion
    logger.info(blue('Adding to notion, let\'s wait for the magic...'))
    try:
      submit_to_notion(cv, desc, amount, list(tags), date)
      logger.info(green('[DONE] ') + blue(
          f'￥{amount} at {date}, {desc} with tags: [{", ".join(tags)}]'))
    except Exception as e:
      logger.error(red('Bad, some error happened: ') + str(e))
      continue

    logger.info(blue('Record submitted, cool! Wanna add more?'))


if __name__ == '__main__':
  print(green('\n\tWelcome to BKP, made by 🐼 with ❤️\n'))

  try:
    main()
  except KeyboardInterrupt:
    print(green('\n\nBYE.'))
    exit(0)