text: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  
  from contextlib import contextmanager
  from dataclasses import dataclass, field
  from typing import Callable, Iterator, List, Optional
  
  from libcst import CSTNode, Module
  from libcst._nodes.internal import CodegenState
  from libcst.metadata.base_provider import BaseMetadataProvider
  
  
  @dataclass(frozen=True)
  class CodeSpan:
      """
      Represents the position of a piece of code by its starting position and length.
  
      Note: This class does not specify the unit of distance - it can be bytes,
      Unicode characters, or something else entirely.
      """
  
      #: Offset of the code from the beginning of the file. Can be 0.
      start: int
      #: Length of the span
      length: int
  
  
  @dataclass(frozen=False)
  class SpanProvidingCodegenState(CodegenState):
      provider: BaseMetadataProvider[CodeSpan]
      get_length: Optional[Callable[[str], int]] = None
      position: int = 0
      _stack: List[int] = field(default_factory=list)
  
      def add_indent_tokens(self) -> None:
          super().add_indent_tokens()
          for token in self.indent_tokens:
              self._update_position(token)
  
      def add_token(self, value: str) -> None:
          super().add_token(value)
          self._update_position(value)
  
      def _update_position(self, value: str) -> None:
          get_length = self.get_length or len
          self.position += get_length(value)
  
      def before_codegen(self, node: CSTNode) -> None:
          self._stack.append(self.position)
  
      def after_codegen(self, node: CSTNode) -> None:
          start = self._stack.pop()
  
          if node not in self.provider._computed:
              end = self.position
              self.provider._computed[node] = CodeSpan(start, length=end - start)
  
      @contextmanager
      def record_syntactic_position(
          self,
          node: CSTNode,
          *,
          start_node: Optional[CSTNode] = None,
          end_node: Optional[CSTNode] = None,
      ) -> Iterator[None]:
          start = self.position
          try:
              yield
          finally:
              end = self.position
              start = (
                  self.provider._computed[start_node].start
                  if start_node is not None
                  else start
              )
              if end_node is not None:
                  end_span = self.provider._computed[end_node]
                  length = (end_span.start + end_span.length) - start
              else:
                  length = end - start
              self.provider._computed[node] = CodeSpan(start, length=length)
  
  
  def byte_length_in_utf8(value: str) -> int:
      return len(value.encode("utf8"))
  
  
  class ByteSpanPositionProvider(BaseMetadataProvider[CodeSpan]):
      """
      Generates offset and length metadata for nodes' positions.
  
      For each :class:`CSTNode` this provider generates a :class:`CodeSpan` that
      contains the byte-offset of the node from the start of the file, and its
      length (also in bytes). The whitespace owned by the node is not included in
      this length.
  
      Note: offset and length measure bytes, not characters (which is significant for
      example in the case of Unicode characters encoded in more than one byte)
      """
  
      def _gen_impl(self, module: Module) -> None:
          state = SpanProvidingCodegenState(
              default_indent=module.default_indent,
              default_newline=module.default_newline,
              provider=self,
              get_length=byte_length_in_utf8,
          )
          module._codegen(state)
  