Skip to content

Instantly share code, notes, and snippets.

@lamchau
Created July 29, 2021 07:36
Show Gist options
  • Save lamchau/7fe42b06b87557e302fb2bce6cf82dae to your computer and use it in GitHub Desktop.
Save lamchau/7fe42b06b87557e302fb2bce6cf82dae to your computer and use it in GitHub Desktop.
from typing import Union, Callable
import types
# 3p libraries
from dateutil.relativedelta import relativedelta # type: ignore
from visidata import Sheet, Column, SettableColumn, asyncthread, vd # type: ignore
import pandas as pd # type: ignore
def get_frequency(column: Column,
f: Union[str, Callable],
column_name: str):
sheet = column.sheet
current_index = sheet.cursorColIndex
frequency_column = SettableColumn(column_name)
sheet.addColumn(frequency_column,
index=current_index + 1)
for r in sheet.rows:
value = r.get(column.name)
if isinstance(f, types.FunctionType):
value = f(pd.Period(value))
elif value:
value = pd.Period(value, freq=f)
frequency_column.setValue(r, value)
@Column.api
@asyncthread
def quarter(column: Column):
get_frequency(column, f='Q', column_name='quarter')
@Column.api
@asyncthread
def year(column: Column):
get_frequency(column, f='Y', column_name='year')
@Column.api
@asyncthread
def week(column: Column):
get_frequency(column, f='W', column_name='week')
@Column.api
@asyncthread
def day(column: Column):
get_frequency(column,
f=lambda x: x.strftime('%A'),
column_name='day')
@Column.api
@asyncthread
def period(column: Column):
get_frequency(column,
f=lambda x: x,
column_name='period')
@Column.api
def diff(column: Column):
sheet = column.sheet
column_names = [x.name for x in sheet.columns]
def to_datetime(value):
return (pd.Period(value)
.to_timestamp()
.to_pydatetime()
.replace(second=0, microsecond=0)
if value else None)
def completer(value, state):
options = [x for x in column_names if x.startswith(value)]
return options[state % len(options)] if options else None
input_column = vd.input('date: ', completer=completer)
if input_column not in column_names:
return None
diff_column = SettableColumn('diff')
diff_seconds_column = SettableColumn('diff_total_seconds', type=int)
index = sheet.cursorColIndex + 1
sheet.addColumn(diff_column, index=index)
sheet.addColumn(diff_seconds_column, index=index)
for r in sheet.rows:
a = to_datetime(r.get(column.name))
b = to_datetime(r.get(input_column))
if a and b:
diff = relativedelta(a, b)
diff_str = f'{diff.years} years(s), {diff.months} month(s), {diff.days} day(s)'
diff_seconds = (a - b).total_seconds()
diff_column.setValue(r, diff_str)
diff_seconds_column.setValue(r, diff_seconds)
# we'll use @ as the prefix since that's the key vd uses to force a date type
Sheet.addCommand(None, '@day', 'cursorCol.day()')
Sheet.addCommand(None, '@diff', 'cursorCol.diff()')
Sheet.addCommand(None, '@period', 'cursorCol.period()')
Sheet.addCommand(None, '@quarter', 'cursorCol.quarter()')
Sheet.addCommand(None, '@week', 'cursorCol.week()')
Sheet.addCommand(None, '@year', 'cursorCol.year()')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment