helpers.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. # Copyright 2022 The HuggingFace Team and Brian Chao. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """
  15. A variety of helper functions and constants when dealing with terminal menu choices, based on
  16. https://github.com/bchao1/bullet
  17. """
  18. import enum
  19. import shutil
  20. import sys
  21. TERMINAL_WIDTH, _ = shutil.get_terminal_size()
  22. CURSOR_TO_CHAR = {"UP": "A", "DOWN": "B", "RIGHT": "C", "LEFT": "D"}
  23. class Direction(enum.Enum):
  24. UP = 0
  25. DOWN = 1
  26. def forceWrite(content, end=""):
  27. sys.stdout.write(str(content) + end)
  28. sys.stdout.flush()
  29. def writeColor(content, color, end=""):
  30. forceWrite(f"\u001b[{color}m{content}\u001b[0m", end)
  31. def reset_cursor():
  32. forceWrite("\r")
  33. def move_cursor(num_lines: int, direction: str):
  34. forceWrite(f"\033[{num_lines}{CURSOR_TO_CHAR[direction.upper()]}")
  35. def clear_line():
  36. forceWrite(" " * TERMINAL_WIDTH)
  37. reset_cursor()
  38. def linebreak():
  39. reset_cursor()
  40. forceWrite("-" * TERMINAL_WIDTH)