Add more generic data-attribute support

Helps solve PostHog/posthog#2140
PostHog · Nov 12, 2020 · e36c9dc · e36c9dc
1 parent 1a8b1f1
commit e36c9dc
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 38 deletions.
diff --git a/__tests__/integration/fixture.js b/__tests__/integration/fixture.js
@@ -60,7 +60,14 @@ export default `
           </div>
       </div>
       <div class="this-is:fine"></div>
-      <div id="somestuff" data-attr="blue"></div>
+      <div class="data-attrs">
+          <div data-attr="blue"></div>
+          <div data-greek-letter="alpha"></div>
+          <div data-greek-letter="omega"></div>
+          <div data-test-attribute-foo>
+              <div data-greek-letter="omega"></div>
+          </div>
+      </div>
       <div>
           <div>
               <div>

diff --git a/__tests__/integration/index.test.js b/__tests__/integration/index.test.js
@@ -22,6 +22,16 @@ const installSimmerOnWindow = windowScope => {
 const createWindow = (dom = '') =>
   installSimmerOnWindow(new JSDOM(`<body>${dom}</body>`).window)
 
+const testExpectSelector = (bodySelector, resultSelector) => {
+  const windowScope = createWindow(fixture)
+  var elements = compareElementsAndSimmer(windowScope, bodySelector)
+  expect(elements).not.toBe(undefined)
+  expect(elements.SimmerEl).not.toBe(undefined)
+  expect(elements.el).not.toBe(undefined)
+  expect(elements.el).toBe(elements.SimmerEl)
+  expect(elements.selector).toEqual(resultSelector)
+}
+
 test(`can analyze an element with an ID`, function () {
   const windowScope = createWindow(fixture)
   var elements = compareElementsAndSimmer(windowScope, '#BodyDiv')
@@ -290,13 +300,16 @@ test(`can analyze an element with a class that contains a colon`, function () {
 })
 
 test(`can analyze an element with data-attr`, function () {
-  const windowScope = createWindow(fixture)
-  var elements = compareElementsAndSimmer(windowScope, '#somestuff')
-  expect(elements).not.toBe(undefined)
-  expect(elements.SimmerEl).not.toBe(undefined)
-  expect(elements.el).not.toBe(undefined)
-  expect(elements.el).toBe(elements.SimmerEl)
-  expect(elements.selector).toEqual("[data-attr='blue']")
+  testExpectSelector('.data-attrs div:nth-child(1)', "[data-attr='blue']")
+})
+
+test(`can analyze an element with data attributes`, function () {
+  testExpectSelector('.data-attrs div:nth-child(2)', "[data-greek-letter='alpha']")
+})
+
+test(`can analyze an element with non-unique data attributes`, function () {
+  // :TODO: This is worse than could actually be generated
+  testExpectSelector('.data-attrs [data-test-attribute-foo] div', "[data-test-attribute-foo] > *")
 })
 
 test(`can't analyze an element which is longer than the selectorMaxLength chars`, function () {

diff --git a/modules/methods/inspectDataAttr.js b/modules/methods/inspectDataAttr.js
@@ -1,5 +1,5 @@
-import { isUniqueDataAttr } from '../queryEngine'
-import { attr } from './validationHelpers'
+import { isUniqueSelector } from '../queryEngine'
+import { validDataSelector } from './validationHelpers'
 /**
  * Inspect the elements' IDs and add them to the CSS Selector
  * @param {array} hierarchy. The hierarchy of elements
@@ -8,31 +8,37 @@ import { attr } from './validationHelpers'
 export default function (hierarchy, state, validateSelector, config, query) {
   return hierarchy.reduce((selectorState, currentElem, index) => {
     if (!selectorState.verified) {
-      const [validatedState] = [currentElem.el.getAttribute('data-attr')]
-        .filter(id => attr(id))
-        .map(dataAttr => {
-          const isUnique = isUniqueDataAttr(query, dataAttr)
-          selectorState.stack[index].push(`[data-attr='${dataAttr}']`)
-          selectorState.specificity += isUnique ? 100 : 50
+      const [bestAttribute] = currentElem.el.getAttributeNames()
+        .filter((attrName) => attrName.indexOf('data-') === 0)
+        .map((attrName) => {
+          const value = currentElem.el.getAttribute(attrName)
+          return value !== '' ? `[${attrName}='${value}']` : `[${attrName}]`
+        })
+        .filter(validDataSelector)
+        .sort()
+        .map((selector) => ({ selector, specificity: isUniqueSelector(query, selector) ? 100 : 50 }))
+        .sort((a, b) => a.specificity - b.specificity)
 
-          if (selectorState.specificity >= config.specificityThreshold) {
-            // we have reached the minimum specificity, lets try verifying now, as this will save us having to add more IDs to the selector
-            if (validateSelector(selectorState)) {
-              // The ID worked like a charm - mark this state as verified and move on!
-              selectorState.verified = true
-            }
-          }
+      if (bestAttribute) {
+        selectorState.stack[index].push(bestAttribute.selector)
+        selectorState.specificity += bestAttribute.specificity
 
-          if (!selectorState.verified && index === 0) {
-            // if the index is 0 then this is the data-attr of the actual element! Which means we have found our selector!
-            // The ID wasn't enough, this means the page, this should never happen as we tested for the data-attr's uniquness, but just incase
-            // we will pop it from the stack as it only adds noise
-            selectorState.stack[index].pop()
-            selectorState.specificity -= isUnique ? 100 : 50
+        if (selectorState.specificity >= config.specificityThreshold) {
+          // we have reached the minimum specificity, lets try verifying now, as this will save us having to add more IDs to the selector
+          if (validateSelector(selectorState)) {
+            // The ID worked like a charm - mark this state as verified and move on!
+            selectorState.verified = true
           }
-          return selectorState
-        })
-      return validatedState || selectorState
+        }
+
+        if (!selectorState.verified && index === 0) {
+          // if the index is 0 then this is the data-attr of the actual element! Which means we have found our selector!
+          // The ID wasn't enough, this means the page, this should never happen as we tested for the data-attr's uniquness, but just incase
+          // we will pop it from the stack as it only adds noise
+          selectorState.stack[index].pop()
+          selectorState.specificity -= bestAttribute.specificity
+        }
+      }
     }
     return selectorState
   }, state)

diff --git a/modules/methods/validationHelpers.js b/modules/methods/validationHelpers.js
@@ -13,6 +13,7 @@ export function tagName (tagName) {
   }
   return false
 }
+
 /**
  * Validate the syntax of an attribute to make sure that it has a valid syntax for the query engine.
  * @param {string} attribute. The element's attribute's value
@@ -27,6 +28,16 @@ export function attr (attribute) {
   return false
 }
 
+export function validDataSelector(selector) {
+  if (
+    typeof selector === 'string' &&
+    selector.match(/^\[data-[0-9a-zA-Z_\-:.]+(='[0-9a-zA-Z_\-:.]*')?\]$/gi) !== null
+  ) {
+    return selector
+  }
+  return false
+}
+
 /**
  * Validate the syntax of an attribute to make sure that it has a valid syntax for the query engine.
  * @param {string} attribute. The element's attribute's value

diff --git a/modules/queryEngine.js b/modules/queryEngine.js
@@ -6,14 +6,11 @@ import { querySelectorAllDeep } from '@mariusandra/query-selector-shadow-dom'
  * @param {object} state. The current selector state (has the stack and specificity sum)
  */
 export function isUniqueElementID (query, elementID) {
-  // use selector to query an element and see if it is a one-to-one selection
-  var results = query(`[id="${elementID}"]`) || []
-  return results.length === 1
+  return isUniqueSelector(query, `[id="${elementID}"]`)
 }
 
-export function isUniqueDataAttr (query, dataAttr) {
-  // use selector to query an element and see if it is a one-to-one selection
-  var results = query(`[data-attr="${dataAttr}"]`) || []
+export function isUniqueSelector(query, selector) {
+  var results = query(selector) || []
   return results.length === 1
 }