Skip to content

Commit

Permalink
piker.config: use tomlkit for accounting files
Browse files Browse the repository at this point in the history
We still need to get some patches landed in order to resolve:
- python-poetry/tomlkit#288
- python-poetry/tomlkit#289
- python-poetry/tomlkit#290

But, this does work for style preservation and the inline-table style we
were previously hacking into the `toml` lib in `.accounting._toml`,
which we can pretty much just drop now B)

Relates to #496 (pretty much solves it near-term i think?)
  • Loading branch information
goodboy committed May 12, 2023
1 parent 5f79434 commit 2865f0e
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 19 deletions.
52 changes: 36 additions & 16 deletions piker/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,13 +23,19 @@
import os
import shutil
import time
from typing import Optional
from typing import (
Callable,
MutableMapping,
)
from pathlib import Path

from bidict import bidict
import toml
import tomli
# import tomlkit # TODO!
import tomlkit
try:
import tomllib
except ModuleNotFoundError:
import tomli as tomllib


from .log import get_logger

Expand Down Expand Up @@ -220,6 +226,11 @@ def load(
conf_name: str = 'brokers',
path: Path | None = None,

decode: Callable[
[str | bytes,],
MutableMapping,
] = tomllib.loads,

**tomlkws,

) -> tuple[dict, Path]:
Expand Down Expand Up @@ -250,11 +261,7 @@ def load(
pass

with path.open(mode='r') as fp:
# TODO: move to tomlkit:
# - needs to be fixed to support bidict?
# - we need to use or fork's fix to do multiline array
# indenting.
config: dict = toml.loads(
config: dict = decode(
fp.read(),
**tomlkws,
)
Expand All @@ -277,7 +284,10 @@ def load_account(
fn: str = f'account.{brokername}.{acctid}.toml'

dirpath: Path = _config_dir / 'accounting'
config, path = load(path=dirpath / fn)
config, path = load(
path=dirpath / fn,
decode=tomlkit.parse,
)

if not config:
legacypath = dirpath / legacy_fn
Expand All @@ -287,14 +297,24 @@ def load_account(
f'Rewriting contents to new name -> {path}\n'
'Please delete the old file!\n'
)
legacy_config, _ = load(path=legacypath)
legacy_config, _ = load(
path=legacypath,

# TODO: move to tomlkit:
# - needs to be fixed to support bidict?
# https://github.com/sdispater/tomlkit/issues/289
# - we need to use or fork's fix to do multiline array
# indenting.
decode=tomlkit.parse,
)
config.update(legacy_config)

# XXX: override the presumably previously non-existant
# file with legacy's contents.
write(
config,
path=path,
fail_empty=False,
)

return config, path
Expand All @@ -321,7 +341,7 @@ def load_ledger(

with fpath.open(mode='rb') as cf:
start = time.time()
ledger_dict = tomli.load(cf)
ledger_dict = tomlkit.parse(cf.read())
log.debug(f'Ledger load took {time.time() - start}s')

return ledger_dict, fpath
Expand Down Expand Up @@ -362,18 +382,18 @@ def write(
f"Writing config `{name}` file to:\n"
f"{path}"
)
with path.open(mode='w') as cf:
return toml.dump(
with path.open(mode='w') as fp:
return tomlkit.dump( # preserve style on write B)
config,
cf,
fp,
**toml_kwargs,
)


def load_accounts(
providers: list[str] | None = None

) -> bidict[str, Optional[str]]:
) -> bidict[str, str | None]:

conf, path = load()
accounts = bidict()
Expand Down
7 changes: 4 additions & 3 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,15 +44,16 @@
]
},
install_requires=[
'toml',
'tomli', # fastest pure py reader
'click',
'tomlkit', # fork & fix for now:
'tomli', # for pre-3.11
'colorlog',
'attrs',
'pygments',
'colorama', # numba traceback coloring
'msgspec', # performant IPC messaging and structs
'protobuf',
'typer',
'rich',

# async
'trio',
Expand Down

0 comments on commit 2865f0e

Please sign in to comment.