BaseSegment.pos_marker is typed as non optional but sometimes set to … (

#2741)

* BaseSegment.pos_marker is typed as non optional but sometimes set to None

* mypy

* Simplify

Co-authored-by: Barry Hart <barry.hart@mailchimp.com>

src/sqlfluff/core/parser/segments/base.py

@@ -84,11 +84,12 @@ class BaseSegment:
     allow_empty = False
     # What other kwargs need to be copied when applying fixes.
     additional_kwargs: List[str] = []
+    pos_marker: Optional[PositionMarker]
 
     def __init__(
         self,
         segments,
-        pos_marker=None,
+        pos_marker: Optional[PositionMarker] = None,
         name: Optional[str] = None,
     ):
         # A cache variable for expandable
@@ -278,6 +279,7 @@ def is_templated(self) -> bool:
         #   Return False for these.
         # * Source string doesn't match raw segment contents. This can only
         #   happen if templating is involved.
+        assert self.pos_marker
         return (
             self.pos_marker.source_slice.start != self.pos_marker.source_slice.stop
             and self.raw != self.pos_marker.source_str()
@@ -1154,6 +1156,7 @@ def iter_patches(self, templated_str: str) -> Iterator[FixPatch]:
         return a hint to deal with that.
         """
         # Does it match? If so we can ignore it.
+        assert self.pos_marker
         matches = self.raw == templated_str[self.pos_marker.templated_slice]
         if matches:
             return

src/sqlfluff/core/rules/base.py

@@ -20,14 +20,14 @@
 import logging
 import pathlib
 import regex
-from typing import Iterable, Optional, List, Set, Tuple, Union, Any
+from typing import cast, Iterable, Optional, List, Set, Tuple, Union, Any
 from collections import namedtuple
 from dataclasses import dataclass
 
 from sqlfluff.core.cached_property import cached_property
 
 from sqlfluff.core.linter import LintedFile
-from sqlfluff.core.parser import BaseSegment, RawSegment
+from sqlfluff.core.parser import BaseSegment, PositionMarker, RawSegment
 from sqlfluff.core.dialects import Dialect
 from sqlfluff.core.errors import SQLLintError
 from sqlfluff.core.rules.functional import Segments
@@ -164,7 +164,7 @@ def __init__(
                         "Developer Note: Edit segment found with preset position "
                         "marker. These should be unset and calculated later."
                     )
-                    seg.pos_marker = None  # type: ignore
+                    seg.pos_marker = None
             # Once stripped, we shouldn't replace any markers because
             # later code may rely on them being accurate, which we
             # can't guarantee with edits.
@@ -269,6 +269,7 @@ def has_template_conflicts(self, templated_file: TemplatedFile) -> bool:
         #    whole anchor segment, because we're not *touching* the anchor
         #    segment, we're inserting **RELATIVE** to it. If ALL are templated,
         #    discard the fix.
+        assert self.anchor.pos_marker
         anchor_slice = self.anchor.pos_marker.templated_slice
         templated_slices = [anchor_slice]
         check_fn = any
@@ -307,7 +308,10 @@ def has_template_conflicts(self, templated_file: TemplatedFile) -> bool:
             return result
 
         # Fix slices were okay. Now check template safety of the "source" field.
-        templated_slices = [source.pos_marker.templated_slice for source in self.source]
+        templated_slices = [
+            cast(PositionMarker, source.pos_marker).templated_slice
+            for source in self.source
+        ]
         raw_slices = self._raw_slices_from_templated_slices(
             templated_file, templated_slices
         )
@@ -736,6 +740,7 @@ def discard_unsafe_fixes(
         fix_slices: Set[RawFileSlice] = set()
         for fix in lint_result.fixes:
             if fix.anchor:
+                assert fix.anchor.pos_marker
                 fix_slices.update(
                     templated_file.raw_slices_spanning_source_slice(
                         fix.anchor.pos_marker.source_slice

src/sqlfluff/core/rules/functional/segment_predicates.py

@@ -142,6 +142,10 @@ def raw_slices(
         raise ValueError(
             'raw_slices: "templated_file" parameter is required.'
         )  # pragma: no cover
+    if not segment.pos_marker:
+        raise ValueError(
+            'raw_slices: "segment" parameter must have pos_marker set.'
+        )  # pragma: no cover
     return RawFileSlices(
         *templated_file.raw_slices_spanning_source_slice(
             segment.pos_marker.source_slice
@@ -159,6 +163,10 @@ def templated_slices(
         raise ValueError(
             'templated_slices: "templated_file" parameter is required.'
         )  # pragma: no cover
+    if not segment.pos_marker:
+        raise ValueError(
+            'templated_slices: "segment" parameter must have pos_marker set.'
+        )  # pragma: no cover
     # :TRICKY: We don't use _find_slice_indices_of_templated_pos() here because
     # it treats TemplatedFileSlice.templated_slice.stop as inclusive, not
     # exclusive. Other parts of SQLFluff rely on this behavior, but we don't
@@ -177,6 +185,10 @@ def templated_slices(
 def raw_slice(segment: BaseSegment, raw_slice_: RawFileSlice) -> str:
     """Return the portion of a segment's source provided by raw_slice."""
     result = ""
+    if not segment.pos_marker:
+        raise ValueError(
+            'raw_slice: "segment" parameter must have pos_marker set.'
+        )  # pragma: no cover
     seg_start = segment.pos_marker.source_slice.start
     seg_stop = segment.pos_marker.source_slice.stop
     if seg_start != seg_stop:

src/sqlfluff/rules/L003.py

@@ -76,6 +76,7 @@ def _segment_length(elem: BaseSegment, tab_space_size: int):
         # whitespace. Check for this.
         if elem.is_type("whitespace") and elem.is_templated:
             # Templated case: Find the leading *literal* whitespace.
+            assert elem.pos_marker
             templated_file = elem.pos_marker.templated_file
             # Extract the leading literal whitespace, slice by slice.
             raw = ""
@@ -855,6 +856,7 @@ def _get_element_template_info(
         if elem.is_type("placeholder"):
             if templated_file is None:
                 raise ValueError("Parameter templated_file cannot be: None.")
+            assert elem.pos_marker
             slices = templated_file.raw_slices_spanning_source_slice(
                 elem.pos_marker.source_slice
             )

src/sqlfluff/rules/L016.py

@@ -399,6 +399,7 @@ def _compute_segment_length(cls, segment: BaseSegment) -> int:
             # line length, and newlines don't make lines longer.
             return 0
 
+        assert segment.pos_marker
         if "\n" in segment.pos_marker.source_str():
             # Similarly we shouldn't see newlines in source segments
             # However for templated loops it's often not possible to
@@ -430,6 +431,7 @@ def _compute_source_length(
             if self.ignore_comment_clauses and segment in memory["comment_clauses"]:
                 continue
 
+            assert segment.pos_marker
             slice = (
                 segment.pos_marker.source_slice.start,
                 segment.pos_marker.source_slice.stop,

src/sqlfluff/rules/L018.py

@@ -1,6 +1,8 @@
 """Implementation of Rule L018."""
 
-from sqlfluff.core.parser import NewlineSegment, WhitespaceSegment
+from typing import cast
+
+from sqlfluff.core.parser import NewlineSegment, PositionMarker, WhitespaceSegment
 
 from sqlfluff.core.rules.base import BaseRule, LintFix, LintResult, RuleContext
 from sqlfluff.core.rules.doc_decorators import (
@@ -115,11 +117,14 @@ def indent_size_up_to(segs):
                             )
                         elif indent_diff > 0:
                             # Is it all whitespace before the bracket on this line?
+                            assert seg.pos_marker
                             prev_segs_on_line = [
                                 elem
                                 for elem in context.segment.raw_segments
-                                if elem.pos_marker.line_no == seg.pos_marker.line_no
-                                and elem.pos_marker.line_pos < seg.pos_marker.line_pos
+                                if cast(PositionMarker, elem.pos_marker).line_no
+                                == seg.pos_marker.line_no
+                                and cast(PositionMarker, elem.pos_marker).line_pos
+                                < seg.pos_marker.line_pos
                             ]
                             if all(
                                 elem.is_type("whitespace") for elem in prev_segs_on_line

src/sqlfluff/rules/L046.py

@@ -53,6 +53,7 @@ def _get_whitespace_ends(s: str) -> Tuple[str, str, str]:
 
     def _eval(self, context: RuleContext) -> LintResult:
         """Look for non-literal segments."""
+        assert context.segment.pos_marker
         if context.segment.is_raw() and not context.segment.pos_marker.is_literal():
             # Does it actually look like a tag?
             templated_raw_slices = context.functional.segment.raw_slices.select(